From 6c2a06b049b5efd9ee6decf3911b5a2cb241230e Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:02:32 +0800 Subject: [PATCH 1/5] refactor(encoding): pass mini-block compression context --- rust/lance-encoding/src/compression.rs | 39 ++++++++++----- .../src/encodings/logical/primitive.rs | 8 +++- .../encodings/logical/primitive/miniblock.rs | 29 +++++++++++- .../logical/primitive/sparse/writer.rs | 6 ++- .../src/encodings/physical/binary.rs | 9 +++- .../src/encodings/physical/bitpacking.rs | 8 +++- .../encodings/physical/byte_stream_split.rs | 20 ++++++-- .../src/encodings/physical/fsst.rs | 10 ++-- .../src/encodings/physical/general.rs | 34 ++++++++++---- .../src/encodings/physical/packed.rs | 10 ++-- .../src/encodings/physical/rle.rs | 47 +++++++++++-------- .../src/encodings/physical/value.rs | 22 ++++++--- 12 files changed, 176 insertions(+), 66 deletions(-) diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 5c15d48c25a..2ac2bedfc75 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -1372,6 +1372,13 @@ mod tests { use arrow_schema::{DataType, Field as ArrowField}; use std::collections::HashMap; + fn miniblock_context() + -> crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext { + crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext::new( + 0, true, true, + ) + } + fn create_test_field(name: &str, data_type: DataType) -> Field { let arrow_field = ArrowField::new(name, data_type, true); let mut field = Field::try_from(&arrow_field).unwrap(); @@ -1759,12 +1766,16 @@ mod tests { let compressor = strategy .create_miniblock_compressor(&field, &fixed_data) .unwrap(); - let (_block, encoding) = compressor.compress(fixed_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(fixed_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, false); let compressor = strategy .create_miniblock_compressor(&field, &variable_data) .unwrap(); - let (_block, encoding) = compressor.compress(variable_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(variable_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, true); // Test pervalue @@ -1794,13 +1805,17 @@ mod tests { let compressor = strategy .create_miniblock_compressor(&field, &fixed_data) .unwrap(); - let (_block, encoding) = compressor.compress(fixed_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(fixed_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, false); let compressor = strategy .create_miniblock_compressor(&field, &variable_data) .unwrap(); - let (_block, encoding) = compressor.compress(variable_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(variable_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, true); // Test pervalue @@ -2097,7 +2112,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2122,7 +2137,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(version); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}"); } } @@ -2150,7 +2165,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2173,7 +2188,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2196,7 +2211,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8); } @@ -2233,7 +2248,7 @@ mod tests { let data = DataBlock::FixedWidth(data); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let rle = expect_rle_encoding(&encoding); assert!( @@ -2281,7 +2296,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { panic!("expected RLE encoding"); }; @@ -2331,7 +2346,7 @@ mod tests { "expected RLE to beat inline bitpacking after child selection, got: {debug_str}" ); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let rle = expect_rle_encoding(&encoding); assert!(matches!( rle.values.as_ref().unwrap().compression.as_ref().unwrap(), diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 1b16ab73b80..cdcb0e0c482 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -55,7 +55,7 @@ use crate::{ encodings::logical::primitive::fullzip::PerValueDataBlock, }; use crate::{ - encodings::logical::primitive::miniblock::MiniBlockCompressed, + encodings::logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, statistics::{ComputeStat, GetStat, Stat}, }; use crate::{ @@ -5282,7 +5282,11 @@ impl PrimitiveStructuralEncoder { let num_items = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; - let (compressed_data, value_encoding) = compressor.compress(data)?; + let common_chunk_buffers = + u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some()); + let compression_context = + MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true); + let (compressed_data, value_encoding) = compressor.compress(data, compression_context)?; let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16; diff --git a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs index edfba526670..9461229310b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs @@ -54,6 +54,29 @@ pub struct MiniBlockCompressed { pub num_values: u64, } +/// Per-page framing details that can affect a mini-block compressor's choice. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MiniBlockCompressionContext { + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, +} + +impl MiniBlockCompressionContext { + /// Creates the framing context supplied by the owning mini-block page. + pub fn new( + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, + ) -> Self { + Self { + common_chunk_buffers, + support_large_chunk, + allow_generic_offsets, + } + } +} + /// Describes the size of a mini-block chunk of data /// /// Mini-block chunks are designed to be small (just a few disk sectors) @@ -113,7 +136,11 @@ pub trait MiniBlockCompressor: std::fmt::Debug + Send + Sync { /// /// This method also returns a description of the encoding applied that will be /// used at decode time to read the data. - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; + fn compress( + &self, + page: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; } #[cfg(test)] diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index 3d82b4a5ef6..809536a2c1b 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -24,7 +24,8 @@ use super::{ SparseValidityMeaning, SparseValiditySet, }; use crate::encodings::logical::primitive::{ - FILL_BYTE, MINIBLOCK_ALIGNMENT, miniblock::MiniBlockCompressed, + FILL_BYTE, MINIBLOCK_ALIGNMENT, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, }; #[derive(Clone, Copy, Default)] @@ -563,7 +564,8 @@ pub fn prepare_values( let num_values = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; - let (compressed, value_compression) = compressor.compress(data)?; + let compression_context = MiniBlockCompressionContext::new(0, support_large_chunk, false); + let (compressed, value_compression) = compressor.compress(data, compression_context)?; let values = serialize_value_chunks(with_explicit_value_counts(compressed)?, support_large_chunk)?; Ok(PreparedSparseValues { diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index d02cf2da693..7c84c2774b7 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -21,7 +21,8 @@ use crate::buffer::LanceBuffer; use crate::data::{BlockInfo, DataBlock, VariableWidthBlock}; use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, + MiniBlockCompressor, }; use crate::format::pb21::CompressiveEncoding; use crate::format::pb21::compressive_encoding::Compression; @@ -245,7 +246,11 @@ impl BinaryMiniBlockEncoder { } impl MiniBlockCompressor for BinaryMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)), _ => Err(Error::invalid_input_source( diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index be0b747e7dc..ad1b567edd7 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -28,7 +28,7 @@ use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompress use crate::data::BlockInfo; use crate::data::{DataBlock, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::pb21::CompressiveEncoding; use crate::format::{ProtobufUtils21, pb21}; @@ -215,7 +215,11 @@ impl InlineBitpacking { } impl MiniBlockCompressor for InlineBitpacking { - fn compress(&self, chunk: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + chunk: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => Ok(self.chunk_data(fixed_width)), _ => Err(Error::invalid_input_source( diff --git a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs index c2b7aac9b9c..b861098d3e4 100644 --- a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs +++ b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs @@ -62,7 +62,7 @@ use crate::compression::MiniBlockDecompressor; use crate::compression_config::BssMode; use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::ProtobufUtils21; use crate::format::pb21::CompressiveEncoding; @@ -107,7 +107,11 @@ impl ByteStreamSplitEncoder { } impl MiniBlockCompressor for ByteStreamSplitEncoder { - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + page: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match page { DataBlock::FixedWidth(data) => { let num_values = data.num_values; @@ -345,7 +349,9 @@ mod tests { }); // Compress - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress let decompressed = decompressor @@ -391,7 +397,9 @@ mod tests { }); // Compress - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress let decompressed = decompressor @@ -424,7 +432,9 @@ mod tests { }); // Compress empty data - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress empty data let decompressed = decompressor.decompress(compressed.data, 0).unwrap(); diff --git a/rust/lance-encoding/src/encodings/physical/fsst.rs b/rust/lance-encoding/src/encodings/physical/fsst.rs index 8c1fe4141df..295ffcd88b1 100644 --- a/rust/lance-encoding/src/encodings/physical/fsst.rs +++ b/rust/lance-encoding/src/encodings/physical/fsst.rs @@ -23,7 +23,7 @@ use crate::{ data::{BlockInfo, DataBlock, VariableWidthBlock}, encodings::logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, }, format::{ ProtobufUtils21, @@ -138,7 +138,11 @@ impl FsstMiniBlockEncoder { } impl MiniBlockCompressor for FsstMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { let compressed = FsstCompressed::fsst_compress(data)?; let data_block = DataBlock::VariableWidth(compressed.data); @@ -148,7 +152,7 @@ impl MiniBlockCompressor for FsstMiniBlockEncoder { as Box; let (binary_miniblock_compressed, binary_array_encoding) = - binary_compressor.compress(data_block)?; + binary_compressor.compress(data_block, context)?; Ok(( binary_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index 53c61928870..769f308ccec 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -9,7 +9,9 @@ use crate::{ compression::MiniBlockDecompressor, data::DataBlock, encodings::{ - logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + logical::primitive::miniblock::{ + MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, + }, physical::block::{CompressionConfig, GeneralBufferCompressor}, }, format::{ProtobufUtils21, pb21::CompressiveEncoding}, @@ -35,9 +37,13 @@ const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024; use super::super::logical::primitive::miniblock::MiniBlockChunk; impl MiniBlockCompressor for GeneralMiniBlockCompressor { - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + page: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { // First, compress with the inner compressor - let (inner_compressed, inner_encoding) = self.inner.compress(page)?; + let (inner_compressed, inner_encoding) = self.inner.compress(page, context)?; // Return the original encoding without compression if there's no data or // the first buffer is not large enough @@ -146,6 +152,10 @@ mod tests { use crate::format::pb21::compressive_encoding::Compression; use arrow_array::{Float64Array, Int32Array}; + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + #[derive(Debug)] struct TestCase { name: &'static str, @@ -249,7 +259,9 @@ mod tests { GeneralMiniBlockCompressor::new(test_case.inner_encoder, test_case.compression); // Compress the data - let (compressed, encoding) = compressor.compress(test_case.data).unwrap(); + let (compressed, encoding) = compressor + .compress(test_case.data, miniblock_context()) + .unwrap(); // Check if compression was applied as expected match &encoding.compression { @@ -461,7 +473,7 @@ mod tests { let compressor = GeneralMiniBlockCompressor::new(inner, compression); // Compress the data - let (compressed, encoding) = compressor.compress(block).unwrap(); + let (compressed, encoding) = compressor.compress(block, miniblock_context()).unwrap(); // Should get GeneralMiniBlock encoding since buffer is 4KB match &encoding.compression { @@ -503,7 +515,7 @@ mod tests { }, ); - let (compressed, _) = compressor.compress(data).unwrap(); + let (compressed, _) = compressor.compress(data, miniblock_context()).unwrap(); // RLE produces 2 buffers, but only the first one is compressed assert_eq!(compressed.data.len(), 2); } @@ -539,7 +551,9 @@ mod tests { }, ); - let (_compressed, encoding) = compressor.compress(test_32.data).unwrap(); + let (_compressed, encoding) = compressor + .compress(test_32.data, miniblock_context()) + .unwrap(); // Verify the encoding structure match &encoding.compression { @@ -596,7 +610,9 @@ mod tests { }, ); - let (_compressed_64, encoding_64) = compressor_64.compress(block_64).unwrap(); + let (_compressed_64, encoding_64) = compressor_64 + .compress(block_64, miniblock_context()) + .unwrap(); // Verify the encoding structure for 64-bit match &encoding_64.compression { @@ -650,7 +666,7 @@ mod tests { }, ); - let result = compressor.compress(empty_block); + let result = compressor.compress(empty_block, miniblock_context()); match result { Ok((compressed, _)) => { assert_eq!(compressed.num_values, 0); diff --git a/rust/lance-encoding/src/encodings/physical/packed.rs b/rust/lance-encoding/src/encodings/physical/packed.rs index ad2221dffed..96aa61e69d6 100644 --- a/rust/lance-encoding/src/encodings/physical/packed.rs +++ b/rust/lance-encoding/src/encodings/physical/packed.rs @@ -27,7 +27,7 @@ use crate::{ }, encodings::logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, }, format::{ ProtobufUtils21, @@ -73,7 +73,11 @@ fn struct_data_block_to_fixed_width_data_block( pub struct PackedStructFixedWidthMiniBlockEncoder {} impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::Struct(struct_data_block) => { let bits_per_values = struct_data_block.children.iter().map(|data_block| data_block.as_fixed_width_ref().unwrap().bits_per_value).collect::>(); @@ -84,7 +88,7 @@ impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { // store and transformed fixed-width data block. let value_miniblock_compressor = Box::new(ValueEncoder::default()) as Box; let (value_miniblock_compressed, value_array_encoding) = - value_miniblock_compressor.compress(data_block)?; + value_miniblock_compressor.compress(data_block, context)?; Ok(( value_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index b04716c2b44..0127817d08e 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -63,7 +63,7 @@ use crate::data::DataBlock; use crate::data::{BlockInfo, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, - MiniBlockCompressor, + MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor}; use crate::format::ProtobufUtils21; @@ -1036,7 +1036,11 @@ impl RleChildCandidate { } impl MiniBlockCompressor for RleEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::FixedWidth(fixed_width) => { let num_values = fixed_width.num_values; @@ -1853,6 +1857,13 @@ mod tests { use arrow_array::Int32Array; use rstest::rstest; + fn compress_miniblock( + compressor: &dyn MiniBlockCompressor, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) + } + fn expand_u16_runs(runs: &RleRuns) -> Vec { let mut expanded = Vec::with_capacity(runs.num_values()); for (value, length) in runs.iter() { @@ -2035,7 +2046,7 @@ mod tests { let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]); let data_block = DataBlock::from_array(array); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, data_block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, data_block).unwrap(); assert_eq!(compressed.num_values, 9); assert_eq!(compressed.chunks.len(), 1); @@ -2056,8 +2067,7 @@ mod tests { data.extend(&[100i32; 300]); // Will be split into 255+45 let array = Int32Array::from(data); - let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Should have 6 runs total (4 for first value, 2 for second) let lengths_buffer = &compressed.data[1]; @@ -2071,7 +2081,7 @@ mod tests { let data = vec![42i32; 1000]; let array = Int32Array::from(data); let (compressed, encoding) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); assert_eq!(compressed.data[0].len(), 4); assert_eq!(compressed.data[1].len(), 2); @@ -2112,7 +2122,7 @@ mod tests { RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false); let array = Int32Array::from(repeating_runs(1024, 4)); let (compressed, encoding) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); let rle = expect_rle(&encoding); assert!(matches!( @@ -2145,7 +2155,7 @@ mod tests { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false); let expected = repeating_runs(1024, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2181,7 +2191,7 @@ mod tests { use crate::encodings::physical::bitpacking::OutOfLineBitpacking; let expected = repeating_runs(1024, 4); - let (compressed, _) = MiniBlockCompressor::compress( + let (compressed, _) = compress_miniblock( &RleEncoder::new(), DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2275,7 +2285,7 @@ mod tests { false, ); let expected = repeating_runs(8192, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2306,7 +2316,7 @@ mod tests { fn test_rle_miniblock_bitpacks_values_child_when_smaller() { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); let expected = monotonic_runs(2048, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2336,7 +2346,7 @@ mod tests { fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); let expected = high_entropy_runs(2048, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2464,7 +2474,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); let decompressor = RleDecompressor::new(bits_per_value); let decompressed = MiniBlockDecompressor::decompress( &decompressor, @@ -2498,7 +2508,7 @@ mod tests { let array = Int32Array::from(data); let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Verify all non-last chunks have power-of-2 values for (i, chunk) in compressed.chunks.iter().enumerate() { @@ -2737,7 +2747,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, empty_block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, empty_block).unwrap(); assert_eq!(compressed.num_values, 0); assert!(compressed.data.is_empty()); @@ -2771,8 +2781,7 @@ mod tests { data.extend(vec![777i32; 2000]); let array = Int32Array::from(data.clone()); - let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Manually decompress all chunks let mut reconstructed = Vec::new(); @@ -2885,7 +2894,7 @@ mod tests { // Compress the data let array = Int32Array::from(data.clone()); let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Decompress and verify match MiniBlockDecompressor::decompress( @@ -2955,7 +2964,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); // Debug first few chunks for (i, chunk) in compressed.chunks.iter().take(5).enumerate() { diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index 606f49b699a..8b7c385f601 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -13,7 +13,7 @@ use crate::data::{ use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, - MiniBlockCompressor, + MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::ProtobufUtils21; use crate::format::pb21::compressive_encoding::Compression; @@ -471,7 +471,11 @@ impl BlockCompressor for ValueEncoder { } impl MiniBlockCompressor for ValueEncoder { - fn compress(&self, chunk: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + chunk: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => { let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); @@ -775,7 +779,7 @@ mod tests { encodings::{ logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::MiniBlockCompressor, + miniblock::{MiniBlockCompressionContext, MiniBlockCompressor}, }, physical::value::ValueDecompressor, }, @@ -789,6 +793,10 @@ mod tests { use super::ValueEncoder; + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + const PRIMITIVE_TYPES: &[DataType] = &[ DataType::Null, DataType::FixedSizeBinary(2), @@ -969,7 +977,8 @@ mod tests { let starting_data = DataBlock::from_array(sample_list.clone()); let encoder = ValueEncoder::default(); - let (data, compression) = MiniBlockCompressor::compress(&encoder, starting_data).unwrap(); + let (data, compression) = + MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); assert_eq!(data.num_values, 3); assert_eq!(data.data.len(), 3); @@ -1030,7 +1039,7 @@ mod tests { let starting_data = DataBlock::from_array(array); let encoder = ValueEncoder::default(); - let result = MiniBlockCompressor::compress(&encoder, starting_data); + let result = MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()); let err = result.expect_err("wide values should not be encodable as miniblock"); assert!( @@ -1142,7 +1151,8 @@ mod tests { ); let encoder = ValueEncoder::default(); - let (data, compression) = MiniBlockCompressor::compress(&encoder, starting_data).unwrap(); + let (data, compression) = + MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { panic!() From 6c7bacaf4c5203be7777141cd3e217fe6db5f0f6 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 16:36:05 +0800 Subject: [PATCH 2/5] refactor(encoding): make block codecs fallible --- rust/lance-encoding/src/compression.rs | 373 ++++++---- rust/lance-encoding/src/compression/block.rs | 80 ++ .../src/compression/block/factory.rs | 533 +++++++++++++ .../src/compression/block/fixed.rs | 124 +++ .../src/compression/block/tests.rs | 289 +++++++ .../src/encodings/logical/primitive.rs | 30 +- .../src/encodings/logical/primitive/sparse.rs | 160 +--- .../logical/primitive/sparse/writer.rs | 66 +- rust/lance-encoding/src/encodings/physical.rs | 32 + .../src/encodings/physical/binary.rs | 56 +- .../src/encodings/physical/bitpacking.rs | 148 +++- .../src/encodings/physical/block.rs | 296 +++++++- .../src/encodings/physical/constant.rs | 144 +++- .../src/encodings/physical/general.rs | 120 ++- .../src/encodings/physical/rle.rs | 703 ++++++++++++++++-- .../src/encodings/physical/value.rs | 119 ++- .../tests/compression_strategy.rs | 81 ++ 17 files changed, 2941 insertions(+), 413 deletions(-) create mode 100644 rust/lance-encoding/src/compression/block.rs create mode 100644 rust/lance-encoding/src/compression/block/factory.rs create mode 100644 rust/lance-encoding/src/compression/block/fixed.rs create mode 100644 rust/lance-encoding/src/compression/block/tests.rs create mode 100644 rust/lance-encoding/tests/compression_strategy.rs diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 2ac2bedfc75..6cb173216b0 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -16,6 +16,10 @@ //! Fullzip compression is a per-value approach where we require that values are transparently //! compressed so that we can locate them later. +pub mod block; + +pub use block::BlockValueType; + #[cfg(feature = "bitpacking")] use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking}; use crate::{ @@ -69,7 +73,9 @@ use crate::{ version::LanceFileVersion, }; -use arrow_array::{cast::AsArray, types::UInt64Type}; +#[cfg(feature = "bitpacking")] +use arrow_array::cast::AsArray; +use arrow_array::types::UInt64Type; use arrow_schema::DataType; use fsst::fsst::{FSST_LEAST_INPUT_MAX_LENGTH, FSST_LEAST_INPUT_SIZE}; use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt}; @@ -99,11 +105,11 @@ const RLE_BLOCK_HEADER_BYTES: u128 = std::mem::size_of::() as u128; /// required (e.g. when encoding metadata buffers like a dictionary or for encoding rep/def /// mini-block chunks) pub trait BlockCompressor: std::fmt::Debug + Send + Sync { - /// Compress the data into a single buffer + /// Compress the data into zero or one buffers. /// - /// Also returns a description of the compression that can be used to decompress - /// when reading the data back - fn compress(&self, data: DataBlock) -> Result; + /// `None` represents a metadata-only codec. `Some` represents a physical + /// payload, including a zero-byte payload. + fn compress(&self, data: DataBlock) -> Result>; } /// A trait to pick which compression to use for given data @@ -119,7 +125,10 @@ pub trait BlockCompressor: std::fmt::Debug + Send + Sync { /// used for narrow data types (both fixed and variable length) where we can /// fit many values into an 16KiB block. pub trait CompressionStrategy: Send + Sync + std::fmt::Debug { - /// Create a block compressor for the given data + /// Create a block compressor for the given data. + /// + /// The returned concrete codec may be reused for independently framed + /// blocks that use the same encoding. fn create_block_compressor( &self, field: &Field, @@ -141,6 +150,18 @@ pub trait CompressionStrategy: Send + Sync + std::fmt::Debug { ) -> Result>; } +pub(crate) fn compress_required_block( + strategy: &dyn CompressionStrategy, + field: &Field, + data: DataBlock, +) -> Result<(LanceBuffer, CompressiveEncoding)> { + let (compressor, encoding) = strategy.create_block_compressor(field, &data)?; + let payload = compressor.compress(data)?.ok_or_else(|| { + Error::internal("Required block compressor selected a metadata-only codec".to_string()) + })?; + Ok((payload, encoding)) +} + #[derive(Debug, Default, Clone)] pub struct DefaultCompressionStrategy { /// User-configured compression parameters @@ -370,8 +391,6 @@ fn try_bitpack_for_mini_block(_data: &FixedWidthDataBlock) -> Option Option { - use arrow_array::cast::AsArray; - let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { return None; @@ -402,6 +421,7 @@ fn estimate_inline_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option { u64::try_from(estimated_bytes).ok() } +#[cfg(feature = "bitpacking")] fn try_bitpack_for_block( data: &FixedWidthDataBlock, ) -> Option<(Box, CompressiveEncoding)> { @@ -435,6 +455,13 @@ fn try_bitpack_for_block( } } +#[cfg(not(feature = "bitpacking"))] +fn try_bitpack_for_block( + _data: &FixedWidthDataBlock, +) -> Option<(Box, CompressiveEncoding)> { + None +} + #[cfg(feature = "bitpacking")] fn estimate_block_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option { let bits = data.bits_per_value; @@ -701,6 +728,63 @@ impl DefaultCompressionStrategy { } impl CompressionStrategy for DefaultCompressionStrategy { + fn create_block_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result<(Box, CompressiveEncoding)> { + let field_params = self.get_merged_field_params(field); + + match data { + DataBlock::FixedWidth(fixed_width) => { + if let Some((compressor, encoding)) = + try_rle_for_block(fixed_width, self.version, &field_params, self.use_rle_v2())? + { + return Ok((compressor, encoding)); + } + if let Some((compressor, encoding)) = try_bitpack_for_block(fixed_width) { + return Ok((compressor, encoding)); + } + + if let Some((compressor, config)) = + try_general_compression(self.version, &field_params, data)? + { + let encoding = ProtobufUtils21::wrapped( + config, + ProtobufUtils21::flat(fixed_width.bits_per_value, None), + )?; + return Ok((compressor, encoding)); + } + + let encoder = Box::new(ValueEncoder::default()); + let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); + Ok((encoder, encoding)) + } + DataBlock::VariableWidth(variable_width) => { + if let Some((compressor, config)) = + try_general_compression(self.version, &field_params, data)? + { + let encoding = ProtobufUtils21::wrapped( + config, + ProtobufUtils21::variable( + ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), + None, + ), + )?; + return Ok((compressor, encoding)); + } + + let encoder = Box::new(VariableEncoder::default()); + let encoding = ProtobufUtils21::variable( + ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), + None, + ); + Ok((encoder, encoding)) + } + _ => unreachable!(), + } + } + fn create_miniblock_compressor( &self, field: &Field, @@ -835,65 +919,6 @@ impl CompressionStrategy for DefaultCompressionStrategy { ), } } - - fn create_block_compressor( - &self, - field: &Field, - data: &DataBlock, - ) -> Result<(Box, CompressiveEncoding)> { - let field_params = self.get_merged_field_params(field); - - match data { - DataBlock::FixedWidth(fixed_width) => { - if let Some((compressor, encoding)) = - try_rle_for_block(fixed_width, self.version, &field_params, self.use_rle_v2())? - { - return Ok((compressor, encoding)); - } - if let Some((compressor, encoding)) = try_bitpack_for_block(fixed_width) { - return Ok((compressor, encoding)); - } - - // Try general compression (user-requested or automatic over MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION) - if let Some((compressor, config)) = - try_general_compression(self.version, &field_params, data)? - { - let encoding = ProtobufUtils21::wrapped( - config, - ProtobufUtils21::flat(fixed_width.bits_per_value, None), - )?; - return Ok((compressor, encoding)); - } - - let encoder = Box::new(ValueEncoder::default()); - let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); - Ok((encoder, encoding)) - } - DataBlock::VariableWidth(variable_width) => { - // Try general compression - if let Some((compressor, config)) = - try_general_compression(self.version, &field_params, data)? - { - let encoding = ProtobufUtils21::wrapped( - config, - ProtobufUtils21::variable( - ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), - None, - ), - )?; - return Ok((compressor, encoding)); - } - - let encoder = Box::new(VariableEncoder::default()); - let encoding = ProtobufUtils21::variable( - ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), - None, - ); - Ok((encoder, encoding)) - } - _ => unreachable!(), - } - } } pub trait MiniBlockDecompressor: std::fmt::Debug + Send + Sync { @@ -915,7 +940,18 @@ pub trait VariablePerValueDecompressor: std::fmt::Debug + Send + Sync { } pub trait BlockDecompressor: std::fmt::Debug + Send + Sync { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result; + fn decompress(&self, data: Option, num_values: u64) -> Result; +} + +pub(crate) fn require_block_payload(data: Option, codec: &str) -> Result { + data.ok_or_else(|| Error::invalid_input(format!("{codec} requires one payload"))) +} + +pub(crate) fn require_no_block_payload(data: Option, codec: &str) -> Result<()> { + if data.is_some() { + return Err(Error::invalid_input(format!("{codec} expects no payload"))); + } + Ok(()) } pub trait DecompressionStrategy: std::fmt::Debug + Send + Sync { @@ -950,7 +986,10 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { description: &CompressiveEncoding, decompression_strategy: &dyn DecompressionStrategy, ) -> Result> { - match description.compression.as_ref().unwrap() { + let compression = description.compression.as_ref().ok_or_else(|| { + Error::invalid_input("Mini-block encoding is missing its compression variant") + })?; + match compression { Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))), #[cfg(feature = "bitpacking")] Compression::InlineBitpacking(description) => { @@ -960,24 +999,14 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Compression::InlineBitpacking(_) => Err(Error::not_supported_source( "this runtime was not built with bitpacking support".into(), )), - Compression::Variable(variable) => { - let Compression::Flat(offsets) = variable - .offsets - .as_ref() - .unwrap() - .compression - .as_ref() - .unwrap() - else { - panic!("Variable compression only supports flat offsets") - }; - Ok(Box::new(BinaryMiniBlockDecompressor::new( - offsets.bits_per_value as u8, - ))) - } + Compression::Variable(variable) => Ok(Box::new( + BinaryMiniBlockDecompressor::from_variable(variable)?, + )), Compression::Fsst(description) => { let inner_decompressor = decompression_strategy.create_miniblock_decompressor( - description.values.as_ref().unwrap(), + description.values.as_ref().ok_or_else(|| { + Error::invalid_input("FSST mini-block is missing its values encoding") + })?, decompression_strategy, )?; Ok(Box::new(FsstMiniBlockDecompressor::new( @@ -1001,10 +1030,18 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { decompression_strategy, )?)), Compression::ByteStreamSplit(bss) => { - let Compression::Flat(values) = - bss.values.as_ref().unwrap().compression.as_ref().unwrap() + let values = bss.values.as_ref().ok_or_else(|| { + Error::invalid_input("ByteStreamSplit is missing its values encoding") + })?; + let Compression::Flat(values) = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input( + "ByteStreamSplit values are missing their compression variant", + ) + })? else { - panic!("ByteStreamSplit compression only supports flat values") + return Err(Error::invalid_input( + "ByteStreamSplit compression only supports flat values", + )); }; Ok(Box::new(ByteStreamSplitDecompressor::new( values.bits_per_value as usize, @@ -1033,7 +1070,13 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { compression_config, ))) } - _ => todo!(), + other => Err(Error::not_supported_source( + format!( + "Mini-block decompression does not support {} encoding", + compression_name(other) + ) + .into(), + )), } } @@ -1128,44 +1171,65 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { &self, description: &CompressiveEncoding, ) -> Result> { - match description.compression.as_ref().unwrap() { - Compression::InlineBitpacking(inline_bitpacking) => Ok(Box::new( - InlineBitpacking::from_description(inline_bitpacking), - )), + let compression = description.compression.as_ref().ok_or_else(|| { + Error::invalid_input("Block encoding is missing its compression variant") + })?; + match compression { Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))), - Compression::Constant(constant) => { - let scalar = constant + Compression::Constant(constant) => Ok(Box::new(ConstantDecompressor::new( + constant .value .as_ref() - .map(|v| LanceBuffer::from_bytes(v.clone(), 1)); - Ok(Box::new(ConstantDecompressor::new(scalar))) - } - Compression::Variable(_) => Ok(Box::new(BinaryBlockDecompressor::default())), - Compression::FixedSizeList(fsl) => { - Ok(Box::new(ValueDecompressor::from_fsl(fsl.as_ref()))) + .map(|value| LanceBuffer::from_bytes(value.clone(), 1)), + ))), + #[cfg(feature = "bitpacking")] + Compression::InlineBitpacking(inline) => { + Ok(Box::new(InlineBitpacking::from_description(inline))) } + #[cfg(not(feature = "bitpacking"))] + Compression::InlineBitpacking(_) => Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )), + #[cfg(feature = "bitpacking")] Compression::OutOfLineBitpacking(out_of_line) => { - // Extract the compressed bit width from the values encoding - let compressed_bit_width = match out_of_line - .values - .as_ref() - .unwrap() - .compression - .as_ref() - .unwrap() - { - Compression::Flat(flat) => flat.bits_per_value, - _ => { - return Err(Error::invalid_input_source( - "OutOfLineBitpacking values must use Flat encoding".into(), - )); - } + let values = out_of_line.values.as_deref().ok_or_else(|| { + Error::invalid_input("OutOfLineBitpacking is missing its values encoding") + })?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input( + "OutOfLineBitpacking values must use Flat encoding", + )); }; Ok(Box::new(OutOfLineBitpacking::new( - compressed_bit_width, + flat.bits_per_value, out_of_line.uncompressed_bits_per_value, ))) } + #[cfg(not(feature = "bitpacking"))] + Compression::OutOfLineBitpacking(_) => Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )), + Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)), + Compression::Variable(variable) => { + let offsets = variable.offsets.as_deref().ok_or_else(|| { + Error::invalid_input("Variable block encoding is missing offsets") + })?; + let Some(Compression::Flat(flat)) = offsets.compression.as_ref() else { + return Err(Error::invalid_input( + "Variable block encoding only supports flat offsets", + )); + }; + if !matches!(flat.bits_per_value, 32 | 64) || flat.data.is_some() { + return Err(Error::invalid_input(format!( + "Variable block offsets require uncompressed 32 or 64-bit Flat encoding, got {} bits", + flat.bits_per_value + ))); + } + Ok(Box::new(BinaryBlockDecompressor::default())) + } + Compression::FixedSizeList(fsl) => { + Ok(Box::new(ValueDecompressor::from_fsl(fsl.as_ref()))) + } Compression::General(general) => { let inner_desc = general .values @@ -1186,8 +1250,13 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Ok(Box::new(general_decompressor)) } - Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)), - _ => todo!(), + other => Err(Error::not_supported_source( + format!( + "Block decompression does not support {} encoding", + compression_name(other) + ) + .into(), + )), } } } @@ -1256,6 +1325,21 @@ fn create_rle_child_decompressor( let (bits_per_value, requires_num_values, needs_decompressor) = validate_rle_child_compression(compression, role)?; + if let Compression::General(general) = compression + && !requires_num_values + { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} general child missing compression config" + )) + })?; + let scheme = compression.scheme().try_into()?; + return RleChildDecompressor::general( + bits_per_value, + CompressionConfig::new(scheme, compression.level), + ); + } + if needs_decompressor { Ok(RleChildDecompressor::block( bits_per_value, @@ -1372,13 +1456,6 @@ mod tests { use arrow_schema::{DataType, Field as ArrowField}; use std::collections::HashMap; - fn miniblock_context() - -> crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext { - crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext::new( - 0, true, true, - ) - } - fn create_test_field(name: &str, data_type: DataType) -> Field { let arrow_field = ArrowField::new(name, data_type, true); let mut field = Field::try_from(&arrow_field).unwrap(); @@ -1386,6 +1463,21 @@ mod tests { field } + fn selected_block_codec( + strategy: &dyn CompressionStrategy, + field: &Field, + data: &DataBlock, + ) -> (Box, CompressiveEncoding) { + strategy.create_block_compressor(field, data).unwrap() + } + + fn miniblock_context() + -> crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext { + crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext::new( + 0, true, true, + ) + } + fn create_fixed_width_block_with_stats( bits_per_value: u64, num_values: u64, @@ -1472,6 +1564,7 @@ mod tests { run_lengths.bits_per_value } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn expect_rle_encoding(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { match encoding.compression.as_ref().unwrap() { Compression::Rle(rle) => rle, @@ -1629,8 +1722,8 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let (compressor, _encoding) = strategy.create_block_compressor(&field, &data).unwrap(); - let debug_str = format!("{:?}", compressor); + let (compressor, _) = selected_block_codec(&strategy, &field, &data); + let debug_str = format!("{compressor:?}"); assert!( debug_str.contains("OutOfLineBitpacking"), "expected OutOfLineBitpacking, got: {debug_str}" @@ -1651,7 +1744,7 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); assert!(format!("{compressor:?}").contains("ValueEncoder")); assert!(matches!( @@ -1679,7 +1772,7 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); let debug_str = format!("{compressor:?}"); assert!( debug_str.contains("OutOfLineBitpacking"), @@ -2538,7 +2631,7 @@ mod tests { let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); let data = create_fixed_width_block(24, 1024); - let (_compressor, encoding) = strategy + let (_, encoding) = strategy .create_block_compressor(&field, &data) .expect("block compressor selection should succeed"); @@ -2569,7 +2662,7 @@ mod tests { "test requires block size above automatic general compression threshold" ); - let (_compressor, encoding) = strategy + let (_, encoding) = strategy .create_block_compressor(&field, &data) .expect("block compressor selection should succeed"); @@ -2594,7 +2687,7 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_3); - let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); assert_eq!(rle_run_length_bits(&encoding), 32); let compressed = compressor.compress(data).unwrap(); @@ -2629,7 +2722,7 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_2); - let (_compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (_, encoding) = selected_block_codec(&strategy, &field, &data); assert_eq!(rle_run_length_bits(&encoding), 8); } @@ -2660,11 +2753,9 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_2); - let (compressor, _) = strategy - .create_block_compressor(&field, &data_block) - .unwrap(); + let (compressor, _) = selected_block_codec(&strategy, &field, &data_block); - let debug_str = format!("{:?}", compressor); + let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); } @@ -2695,11 +2786,9 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_1); - let (compressor, _) = strategy - .create_block_compressor(&field, &data_block) - .unwrap(); + let (compressor, _) = selected_block_codec(&strategy, &field, &data_block); - let debug_str = format!("{:?}", compressor); + let debug_str = format!("{compressor:?}"); assert!( !debug_str.contains("RleEncoder"), "RLE should not be used for V2.1" diff --git a/rust/lance-encoding/src/compression/block.rs b/rust/lance-encoding/src/compression/block.rs new file mode 100644 index 00000000000..4ca9e9785db --- /dev/null +++ b/rust/lance-encoding/src/compression/block.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Shared types and decoder construction for unsigned block sequences. +//! +//! As with mini-block compression, selectors return concrete codecs. Codecs +//! own their child codecs, framing, and validation. + +use lance_core::{Error, Result}; + +#[cfg(feature = "bitpacking")] +pub(crate) const BITPACK_CHUNK_VALUES: u64 = 1024; + +/// Typed unsigned role expected from a block descriptor. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BlockValueType { + /// Unsigned 8-bit values. + UInt8, + /// Unsigned 16-bit values. + UInt16, + /// Unsigned 32-bit values. + UInt32, + /// Unsigned 64-bit values. + UInt64, +} + +impl BlockValueType { + pub(crate) fn from_bits(bits_per_value: u64) -> Result { + match bits_per_value { + 8 => Ok(Self::UInt8), + 16 => Ok(Self::UInt16), + 32 => Ok(Self::UInt32), + 64 => Ok(Self::UInt64), + _ => Err(Error::invalid_input(format!( + "Block sequence only supports 8, 16, 32, or 64-bit values, got {bits_per_value}" + ))), + } + } + + /// Returns the fixed width of each value. + pub fn bits_per_value(self) -> u64 { + match self { + Self::UInt8 => 8, + Self::UInt16 => 16, + Self::UInt32 => 32, + Self::UInt64 => 64, + } + } + + pub(crate) fn bytes_per_value(self) -> usize { + (self.bits_per_value() / 8) as usize + } + + #[cfg(test)] + pub(crate) fn max_value(self) -> u64 { + match self { + Self::UInt8 => u8::MAX as u64, + Self::UInt16 => u16::MAX as u64, + Self::UInt32 => u32::MAX as u64, + Self::UInt64 => u64::MAX, + } + } +} + +mod factory; +pub(crate) mod fixed; + +#[cfg(test)] +pub(crate) use factory::encode_scalar; +pub(crate) use factory::{create_block_decompressor, validate_fixed_payload_len}; +#[cfg(feature = "bitpacking")] +pub(crate) use factory::{validate_inline_bitpacking_payload, validate_out_of_line_payload}; +#[cfg(test)] +pub(crate) use fixed::fixed_from_u64_values; +#[cfg(any(test, feature = "bitpacking"))] +pub(crate) use fixed::visit_unsigned_values; +pub(crate) use fixed::{fixed_block, read_unsigned_values}; + +#[cfg(test)] +mod tests; diff --git a/rust/lance-encoding/src/compression/block/factory.rs b/rust/lance-encoding/src/compression/block/factory.rs new file mode 100644 index 00000000000..55bd03376d7 --- /dev/null +++ b/rust/lance-encoding/src/compression/block/factory.rs @@ -0,0 +1,533 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Fallible construction of concrete generic block decompressors. + +#[cfg(test)] +use bytes::Bytes; + +use super::*; +use crate::{ + buffer::LanceBuffer, + compression::BlockDecompressor, + encodings::physical::{ + block::{CompressionConfig, CompressionScheme}, + constant::ConstantBlockDecompressor, + general::GenericGeneralBlockDecompressor, + rle::{BlockRleDecompressor, BlockRunCount, MetadataRunLengths}, + value::FixedWidthBlockDecompressor, + }, + format::pb21::{self, CompressiveEncoding, compressive_encoding::Compression}, +}; + +#[cfg(feature = "bitpacking")] +use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking}; + +fn checked_payload_bytes( + value_type: BlockValueType, + num_values: u64, + label: &str, +) -> Result { + let bytes = usize::try_from(num_values) + .ok() + .and_then(|num_values| num_values.checked_mul(value_type.bytes_per_value())) + .ok_or_else(|| Error::invalid_input(format!("{label} payload length overflows usize")))?; + if bytes > isize::MAX as usize { + return Err(Error::invalid_input(format!( + "{label} payload length {bytes} exceeds isize::MAX" + ))); + } + Ok(bytes) +} + +pub fn validate_fixed_payload_len( + payload: &LanceBuffer, + value_type: BlockValueType, + num_values: u64, + label: &str, +) -> Result<()> { + let expected = checked_payload_bytes(value_type, num_values, label)?; + if payload.len() != expected { + return Err(Error::invalid_input(format!( + "{label} payload has {} bytes, expected {expected} for {num_values} {}-bit values", + payload.len(), + value_type.bits_per_value() + ))); + } + Ok(()) +} + +#[cfg(feature = "bitpacking")] +pub fn validate_inline_bitpacking_payload( + payload: &LanceBuffer, + value_type: BlockValueType, + num_values: u64, +) -> Result<()> { + if num_values == 0 || num_values > BITPACK_CHUNK_VALUES { + return Err(Error::invalid_input(format!( + "Inline bitpacking cardinality {num_values} is outside 1..={BITPACK_CHUNK_VALUES}" + ))); + } + let word_bytes = value_type.bytes_per_value(); + if payload.len() < word_bytes { + return Err(Error::invalid_input(format!( + "Inline bitpacking payload has {} bytes, shorter than its {word_bytes}-byte header", + payload.len() + ))); + } + let bit_width = decode_scalar(&payload[..word_bytes], value_type, "Inline bitpacking")?; + if bit_width > value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Inline bitpacking width {bit_width} exceeds {}", + value_type.bits_per_value() + ))); + } + let packed_words = (BITPACK_CHUNK_VALUES * bit_width) / value_type.bits_per_value(); + let expected_words = 1_u64 + .checked_add(packed_words) + .ok_or_else(|| Error::invalid_input("Inline bitpacking payload word count overflows"))?; + let expected_bytes = usize::try_from(expected_words) + .ok() + .and_then(|words| words.checked_mul(word_bytes)) + .ok_or_else(|| Error::invalid_input("Inline bitpacking payload length overflows"))?; + if payload.len() != expected_bytes { + return Err(Error::invalid_input(format!( + "Inline bitpacking payload has {} bytes, expected {expected_bytes}", + payload.len() + ))); + } + Ok(()) +} + +#[cfg(feature = "bitpacking")] +fn out_of_line_payload_bytes( + value_type: BlockValueType, + num_values: u64, + compressed_bits_per_value: u64, +) -> Result { + if compressed_bits_per_value >= value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Invalid out-of-line bit width {compressed_bits_per_value} for {}-bit values", + value_type.bits_per_value() + ))); + } + let full_chunks = num_values / BITPACK_CHUNK_VALUES; + let tail_values = num_values % BITPACK_CHUNK_VALUES; + let words_per_chunk = + (BITPACK_CHUNK_VALUES * compressed_bits_per_value).div_ceil(value_type.bits_per_value()); + let mut words = full_chunks + .checked_mul(words_per_chunk) + .ok_or_else(|| Error::invalid_input("Out-of-line bitpacking word count overflows"))?; + if tail_values > 0 { + let tail_bit_savings = value_type.bits_per_value() - compressed_bits_per_value; + let padding_cost = compressed_bits_per_value * (BITPACK_CHUNK_VALUES - tail_values); + let tail_pack_savings = tail_bit_savings * tail_values; + words = words + .checked_add(if padding_cost < tail_pack_savings { + words_per_chunk + } else { + tail_values + }) + .ok_or_else(|| { + Error::invalid_input("Out-of-line bitpacking tail word count overflows") + })?; + } + words + .checked_mul(value_type.bytes_per_value() as u64) + .ok_or_else(|| Error::invalid_input("Out-of-line bitpacking byte length overflows")) +} + +#[cfg(feature = "bitpacking")] +pub fn validate_out_of_line_payload( + payload: &LanceBuffer, + value_type: BlockValueType, + num_values: u64, + compressed_bits_per_value: u64, +) -> Result<()> { + let expected = out_of_line_payload_bytes(value_type, num_values, compressed_bits_per_value)?; + if payload.len() as u64 != expected { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking payload has {} bytes, expected {expected}", + payload.len() + ))); + } + Ok(()) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Position { + Root, + Child, +} + +/// Builds the concrete decoder tree while validating the bounded block grammar. +pub fn create_block_decompressor( + encoding: &CompressiveEncoding, + expected_type: BlockValueType, +) -> Result<(Box, bool)> { + create_inner(encoding, expected_type, Position::Root, true) +} + +fn create_inner( + encoding: &CompressiveEncoding, + expected_type: BlockValueType, + position: Position, + allow_general: bool, +) -> Result<(Box, bool)> { + let compression = encoding + .compression + .as_ref() + .ok_or_else(|| Error::invalid_input("Block encoding is missing its compression variant"))?; + let expected_bits = expected_type.bits_per_value(); + match compression { + Compression::Flat(flat) => { + validate_flat(flat, expected_type, "Flat")?; + Ok(( + Box::new(FixedWidthBlockDecompressor::new(expected_type)), + true, + )) + } + Compression::Constant(constant) => { + let value = constant + .value + .as_deref() + .map(|value| decode_scalar(value, expected_type, "Constant")) + .transpose()?; + Ok(( + Box::new(ConstantBlockDecompressor::new(expected_type, value)), + false, + )) + } + Compression::InlineBitpacking(bitpacking) => { + validate_declared_bits( + bitpacking.uncompressed_bits_per_value, + expected_type, + "Inline bitpacking", + )?; + if bitpacking.values.is_some() { + return Err(Error::invalid_input( + "Inline bitpacking leaf buffer compression is unsupported", + )); + } + #[cfg(feature = "bitpacking")] + { + Ok((Box::new(InlineBitpacking::new(expected_bits)), true)) + } + #[cfg(not(feature = "bitpacking"))] + { + Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )) + } + } + Compression::OutOfLineBitpacking(bitpacking) => { + validate_declared_bits( + bitpacking.uncompressed_bits_per_value, + expected_type, + "Out-of-line bitpacking", + )?; + let values = bitpacking.values.as_deref().ok_or_else(|| { + Error::invalid_input("Out-of-line bitpacking is missing its values encoding") + })?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input( + "Out-of-line bitpacking values must use Flat encoding", + )); + }; + if flat.data.is_some() || flat.bits_per_value >= expected_bits { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking width {} must be between 0 and {}", + flat.bits_per_value, + expected_bits - 1 + ))); + } + #[cfg(feature = "bitpacking")] + { + Ok(( + Box::new(OutOfLineBitpacking::new(flat.bits_per_value, expected_bits)), + true, + )) + } + #[cfg(not(feature = "bitpacking"))] + { + Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )) + } + } + Compression::General(general) => { + if position != Position::Root || !allow_general { + return Err(Error::invalid_input( + "General compression is only supported as the single outer block transform", + )); + } + let config = validate_compression_config(general.compression.as_ref(), "General")?; + let child_encoding = general.values.as_deref().ok_or_else(|| { + Error::invalid_input("General compression is missing its child encoding") + })?; + if !matches!( + child_encoding.compression.as_ref(), + Some(Compression::Flat(_)) + ) { + return Err(Error::invalid_input( + "Outer General block compression only supports a Flat child", + )); + } + let (child, child_has_payload) = + create_inner(child_encoding, expected_type, Position::Child, false)?; + if !child_has_payload { + return Err(Error::invalid_input( + "Outer General block compression requires a payload-bearing child", + )); + } + Ok(( + Box::new(GenericGeneralBlockDecompressor::new( + child, + config, + expected_type, + )), + true, + )) + } + Compression::Rle(rle) => { + if position != Position::Root { + return Err(Error::invalid_input( + "RLE is not supported as a block codec child", + )); + } + let values_encoding = rle + .values + .as_deref() + .ok_or_else(|| Error::invalid_input("RLE is missing its values encoding"))?; + let run_lengths_encoding = rle + .run_lengths + .as_deref() + .ok_or_else(|| Error::invalid_input("RLE is missing its run lengths encoding"))?; + let run_length_type = infer_inner(run_lengths_encoding, Position::Child)?; + if !matches!( + run_length_type, + BlockValueType::UInt8 | BlockValueType::UInt16 | BlockValueType::UInt32 + ) { + return Err(Error::invalid_input(format!( + "RLE run lengths must use 8, 16, or 32-bit values, got {}", + run_length_type.bits_per_value() + ))); + } + let (values, values_have_payload) = + create_inner(values_encoding, expected_type, Position::Child, false)?; + let (run_lengths, run_lengths_have_payload) = create_inner( + run_lengths_encoding, + run_length_type, + Position::Child, + false, + )?; + let metadata_run_lengths = match run_lengths_encoding.compression.as_ref() { + Some(Compression::Constant(constant)) => { + let value = decode_scalar( + constant.value.as_deref().ok_or_else(|| { + Error::invalid_input("RLE run lengths Constant is missing its scalar") + })?, + run_length_type, + "RLE run lengths Constant", + )?; + Some(MetadataRunLengths::Constant(value)) + } + _ => None, + }; + let run_count = if let Some(metadata) = metadata_run_lengths { + BlockRunCount::Metadata(metadata) + } else if matches!( + run_lengths_encoding.compression.as_ref(), + Some(Compression::Flat(_)) + ) { + BlockRunCount::RunLengthsPayload + } else if matches!( + values_encoding.compression.as_ref(), + Some(Compression::Flat(_)) + ) { + BlockRunCount::ValuesPayload + } else { + return Err(Error::invalid_input( + "RLE requires metadata run lengths or a Flat child to determine the run count", + )); + }; + Ok(( + Box::new(BlockRleDecompressor::new( + expected_type, + run_length_type, + values, + run_lengths, + values_have_payload, + run_lengths_have_payload, + run_count, + )), + values_have_payload || run_lengths_have_payload, + )) + } + other => Err(Error::invalid_input(format!( + "Unsupported block sequence encoding: {}", + compression_name(other) + ))), + } +} + +fn validate_flat(flat: &pb21::Flat, expected_type: BlockValueType, label: &str) -> Result<()> { + validate_declared_bits(flat.bits_per_value, expected_type, label)?; + if flat.data.is_some() { + return Err(Error::invalid_input(format!( + "{label} leaf buffer compression is unsupported" + ))); + } + Ok(()) +} + +fn validate_declared_bits(actual: u64, expected_type: BlockValueType, label: &str) -> Result<()> { + if actual != expected_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "{label} declares {actual}-bit values, expected {}", + expected_type.bits_per_value() + ))); + } + Ok(()) +} + +fn decode_scalar(bytes: &[u8], value_type: BlockValueType, label: &str) -> Result { + if bytes.len() != value_type.bytes_per_value() { + return Err(Error::invalid_input(format!( + "{label} scalar has {} bytes, expected {}", + bytes.len(), + value_type.bytes_per_value() + ))); + } + Ok(match value_type { + BlockValueType::UInt8 => u64::from(bytes[0]), + BlockValueType::UInt16 => u64::from(u16::from_le_bytes([bytes[0], bytes[1]])), + BlockValueType::UInt32 => u64::from(u32::from_le_bytes( + bytes.try_into().expect("scalar length was checked"), + )), + BlockValueType::UInt64 => { + u64::from_le_bytes(bytes.try_into().expect("scalar length was checked")) + } + }) +} + +#[cfg(test)] +pub fn encode_scalar(value: u64, value_type: BlockValueType) -> Result { + if value > value_type.max_value() { + return Err(Error::invalid_input(format!( + "Scalar value {value} exceeds the {}-bit value range", + value_type.bits_per_value() + ))); + } + Ok(match value_type { + BlockValueType::UInt8 => vec![value as u8], + BlockValueType::UInt16 => (value as u16).to_le_bytes().to_vec(), + BlockValueType::UInt32 => (value as u32).to_le_bytes().to_vec(), + BlockValueType::UInt64 => value.to_le_bytes().to_vec(), + } + .into()) +} + +fn validate_compression_config( + compression: Option<&pb21::BufferCompression>, + label: &str, +) -> Result { + let compression = compression + .ok_or_else(|| Error::invalid_input(format!("{label} is missing compression config")))?; + let scheme = pb21::CompressionScheme::try_from(compression.scheme).map_err(|_| { + Error::invalid_input(format!( + "{label} has unknown compression scheme {}", + compression.scheme + )) + })?; + let scheme = CompressionScheme::try_from(scheme)?; + Ok(CompressionConfig::new(scheme, compression.level)) +} + +fn infer_inner(encoding: &CompressiveEncoding, position: Position) -> Result { + let compression = encoding + .compression + .as_ref() + .ok_or_else(|| Error::invalid_input("Block encoding is missing its compression variant"))?; + match compression { + Compression::Flat(flat) => BlockValueType::from_bits(flat.bits_per_value), + Compression::Constant(constant) => { + let value = constant.value.as_ref().ok_or_else(|| { + Error::invalid_input( + "Cannot infer an empty Constant block type without its typed container role", + ) + })?; + BlockValueType::from_bits((value.len() * 8) as u64) + } + Compression::InlineBitpacking(bitpacking) => { + BlockValueType::from_bits(bitpacking.uncompressed_bits_per_value) + } + Compression::OutOfLineBitpacking(bitpacking) => { + BlockValueType::from_bits(bitpacking.uncompressed_bits_per_value) + } + Compression::General(general) if position == Position::Root => infer_inner( + general + .values + .as_deref() + .ok_or_else(|| Error::invalid_input("General is missing its child encoding"))?, + Position::Child, + ), + Compression::Rle(rle) if position == Position::Root => infer_inner( + rle.values + .as_deref() + .ok_or_else(|| Error::invalid_input("RLE is missing its values encoding"))?, + Position::Child, + ), + other => Err(Error::invalid_input(format!( + "Cannot infer bounded block value type from {} at {position:?}", + compression_name(other) + ))), + } +} + +fn compression_name(compression: &Compression) -> &'static str { + match compression { + Compression::Flat(_) => "flat", + Compression::Variable(_) => "variable", + Compression::Constant(_) => "constant", + Compression::OutOfLineBitpacking(_) => "out-of-line bitpacking", + Compression::InlineBitpacking(_) => "inline bitpacking", + Compression::Fsst(_) => "fsst", + Compression::Dictionary(_) => "dictionary", + Compression::Rle(_) => "rle", + Compression::ByteStreamSplit(_) => "byte stream split", + Compression::General(_) => "general", + Compression::FixedSizeList(_) => "fixed-size list", + Compression::PackedStruct(_) => "packed struct", + Compression::VariablePackedStruct(_) => "variable packed struct", + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_nested_transform_trees() { + let nested = crate::format::ProtobufUtils21::rle( + crate::format::ProtobufUtils21::rle( + crate::format::ProtobufUtils21::flat(64, None), + crate::format::ProtobufUtils21::constant(Some(vec![1].into())), + ), + crate::format::ProtobufUtils21::constant(Some(vec![1].into())), + ); + assert!( + create_block_decompressor(&nested, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("child") + ); + } + + #[test] + fn rejects_mistyped_flat_leaf() { + let flat = crate::format::ProtobufUtils21::flat(32, None); + let error = create_block_decompressor(&flat, BlockValueType::UInt64).unwrap_err(); + assert!(error.to_string().contains("expected 64")); + } +} diff --git a/rust/lance-encoding/src/compression/block/fixed.rs b/rust/lance-encoding/src/compression/block/fixed.rs new file mode 100644 index 00000000000..d5288d571ef --- /dev/null +++ b/rust/lance-encoding/src/compression/block/fixed.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use super::*; +use crate::{ + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, +}; + +pub fn fixed_block(value_type: BlockValueType, num_values: u64, data: LanceBuffer) -> DataBlock { + DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: value_type.bits_per_value(), + data, + num_values, + block_info: BlockInfo::default(), + }) +} + +#[cfg(test)] +pub fn fixed_from_u64_values( + values: &[u64], + value_type: BlockValueType, + label: &str, +) -> Result { + let data = match value_type { + BlockValueType::UInt8 => LanceBuffer::reinterpret_vec( + values + .iter() + .map(|value| { + u8::try_from(*value).map_err(|_| { + Error::invalid_input(format!("{label} value {value} exceeds u8::MAX")) + }) + }) + .collect::>>()?, + ), + BlockValueType::UInt16 => LanceBuffer::reinterpret_vec( + values + .iter() + .map(|value| { + u16::try_from(*value).map_err(|_| { + Error::invalid_input(format!("{label} value {value} exceeds u16::MAX")) + }) + }) + .collect::>>()?, + ), + BlockValueType::UInt32 => LanceBuffer::reinterpret_vec( + values + .iter() + .map(|value| { + u32::try_from(*value).map_err(|_| { + Error::invalid_input(format!("{label} value {value} exceeds u32::MAX")) + }) + }) + .collect::>>()?, + ), + BlockValueType::UInt64 => LanceBuffer::reinterpret_vec(values.to_vec()), + }; + Ok(FixedWidthDataBlock { + bits_per_value: value_type.bits_per_value(), + data, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + }) +} + +#[cfg(any(test, feature = "bitpacking"))] +pub fn visit_unsigned_values( + block: &FixedWidthDataBlock, + value_type: BlockValueType, + mut visit: impl FnMut(u64) -> Result<()>, +) -> Result<()> { + validate_fixed_payload_len(&block.data, value_type, block.num_values, "Block input")?; + match value_type { + BlockValueType::UInt8 => { + for value in block.data.iter().copied() { + visit(u64::from(value))?; + } + } + BlockValueType::UInt16 => { + for value in block.data.borrow_to_typed_view::().iter().copied() { + visit(u64::from(value))?; + } + } + BlockValueType::UInt32 => { + for value in block.data.borrow_to_typed_view::().iter().copied() { + visit(u64::from(value))?; + } + } + BlockValueType::UInt64 => { + for value in block.data.borrow_to_typed_view::().iter().copied() { + visit(value)?; + } + } + } + Ok(()) +} + +pub fn read_unsigned_values( + block: &FixedWidthDataBlock, + value_type: BlockValueType, +) -> Result> { + validate_fixed_payload_len( + &block.data, + value_type, + block.num_values, + "Fixed-width block", + )?; + Ok(match value_type { + BlockValueType::UInt8 => block.data.iter().map(|value| u64::from(*value)).collect(), + BlockValueType::UInt16 => block + .data + .borrow_to_typed_slice::() + .iter() + .map(|value| u64::from(*value)) + .collect(), + BlockValueType::UInt32 => block + .data + .borrow_to_typed_slice::() + .iter() + .map(|value| u64::from(*value)) + .collect(), + BlockValueType::UInt64 => block.data.borrow_to_typed_slice::().to_vec(), + }) +} diff --git a/rust/lance-encoding/src/compression/block/tests.rs b/rust/lance-encoding/src/compression/block/tests.rs new file mode 100644 index 00000000000..372cf56a2b1 --- /dev/null +++ b/rust/lance-encoding/src/compression/block/tests.rs @@ -0,0 +1,289 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use super::*; + +use crate::{ + buffer::LanceBuffer, + compression::BlockCompressor, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + encodings::physical::{ + constant::ConstantBlockCompressor, rle::BlockRleCompressor, + value::FixedWidthBlockCompressor, + }, + format::{ProtobufUtils21, pb21::CompressiveEncoding}, +}; + +fn fixed_u64(values: &[u64]) -> FixedWidthDataBlock { + FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(values.to_vec()), + bits_per_value: 64, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + } +} + +fn decoded_u64(block: DataBlock) -> Vec { + let DataBlock::FixedWidth(block) = block else { + panic!("expected fixed-width output"); + }; + block.data.borrow_to_typed_slice::().to_vec() +} + +#[test] +fn scalar_encoding_is_little_endian_and_bounded() { + assert_eq!( + encode_scalar(0x0102, BlockValueType::UInt16) + .unwrap() + .as_ref(), + &[0x02, 0x01] + ); + assert!(encode_scalar(256, BlockValueType::UInt8).is_err()); +} + +fn round_trip_u64( + values: &[u64], + compressor: Box, + encoding: CompressiveEncoding, +) { + let payload = compressor + .compress(DataBlock::FixedWidth(fixed_u64(values))) + .unwrap(); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert_eq!(payload.is_some(), has_payload); + assert_eq!( + decoded_u64(decoder.decompress(payload, values.len() as u64).unwrap()), + values + ); +} + +#[test] +fn metadata_compressors_validate_and_decode() { + round_trip_u64( + &[], + Box::new(ConstantBlockCompressor::new(BlockValueType::UInt64, None)), + ProtobufUtils21::constant(None), + ); + round_trip_u64( + &[7_u64; 32], + Box::new(ConstantBlockCompressor::new( + BlockValueType::UInt64, + Some(7), + )), + ProtobufUtils21::constant(Some(7_u64.to_le_bytes().to_vec().into())), + ); +} + +#[test] +fn payload_presence_distinguishes_metadata_from_empty_payload() { + let flat = Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)); + let payload = flat + .compress(DataBlock::FixedWidth(fixed_u64(&[]))) + .unwrap(); + assert!(payload.as_ref().is_some_and(|payload| payload.is_empty())); + + let (flat_decoder, flat_has_payload) = + create_block_decompressor(&ProtobufUtils21::flat(64, None), BlockValueType::UInt64) + .unwrap(); + assert!(flat_has_payload); + assert!(flat_decoder.decompress(None, 0).is_err()); + assert!( + decoded_u64( + flat_decoder + .decompress(Some(LanceBuffer::empty()), 0) + .unwrap() + ) + .is_empty() + ); + + let (constant_decoder, constant_has_payload) = create_block_decompressor( + &ProtobufUtils21::constant(Some(7_u64.to_le_bytes().to_vec().into())), + BlockValueType::UInt64, + ) + .unwrap(); + assert!(!constant_has_payload); + assert!( + constant_decoder + .decompress(Some(LanceBuffer::empty()), 1) + .is_err() + ); +} + +#[test] +fn rle_compressor_owns_and_reuses_children() { + let mut values = Vec::new(); + for value in [11_u64, 91, 37, 123] { + values.extend(std::iter::repeat_n(value, 256)); + } + let compressor = Box::new(BlockRleCompressor::new( + BlockValueType::UInt64, + BlockValueType::UInt16, + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)), + Box::new(ConstantBlockCompressor::new( + BlockValueType::UInt16, + Some(256), + )), + )); + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(64, None), + ProtobufUtils21::constant(Some(256_u16.to_le_bytes().to_vec().into())), + ); + round_trip_u64(&values, compressor, encoding); +} + +#[cfg(feature = "bitpacking")] +#[test] +fn out_of_line_bitpacking_round_trip() { + use crate::encodings::physical::bitpacking::OutOfLineBitpacking; + + let values = (0..4096_u64) + .map(|index| (index * 37) % 1024) + .collect::>(); + round_trip_u64( + &values, + Box::new(OutOfLineBitpacking::new(10, 64)), + ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(10, None)), + ); +} + +#[cfg(any(feature = "lz4", feature = "zstd"))] +#[test] +fn general_compressor_owns_a_flat_child() { + use crate::encodings::physical::{ + block::{CompressionConfig, CompressionScheme}, + general::GeneralBlockCompressor, + }; + + let scheme = if cfg!(feature = "lz4") { + CompressionScheme::Lz4 + } else { + CompressionScheme::Zstd + }; + let config = CompressionConfig::new(scheme, None); + let compressor = Box::new(GeneralBlockCompressor::new( + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)), + config, + )); + let encoding = ProtobufUtils21::wrapped(config, ProtobufUtils21::flat(64, None)).unwrap(); + let values = (0..16_384_u64).map(|value| value % 17).collect::>(); + round_trip_u64(&values, compressor, encoding); +} + +#[test] +fn factory_rejects_unbounded_or_mistyped_trees() { + let nested_rle = ProtobufUtils21::rle( + ProtobufUtils21::rle( + ProtobufUtils21::flat(64, None), + ProtobufUtils21::constant(Some(vec![1].into())), + ), + ProtobufUtils21::constant(Some(vec![1].into())), + ); + assert!( + create_block_decompressor(&nested_rle, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("child") + ); + + let wrong_width = ProtobufUtils21::flat(32, None); + assert!( + create_block_decompressor(&wrong_width, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("expected 64") + ); +} + +#[test] +fn constant_cardinality_contract_is_checked_at_decode() { + let (empty, has_payload) = + create_block_decompressor(&ProtobufUtils21::constant(None), BlockValueType::UInt64) + .unwrap(); + assert!(!has_payload); + assert!(decoded_u64(empty.decompress(None, 0).unwrap()).is_empty()); + assert!(empty.decompress(None, 1).is_err()); + + let (present, has_payload) = create_block_decompressor( + &ProtobufUtils21::constant(Some(7_u64.to_le_bytes().to_vec().into())), + BlockValueType::UInt64, + ) + .unwrap(); + assert!(!has_payload); + assert!(present.decompress(None, 0).is_err()); + assert_eq!( + decoded_u64(present.decompress(None, 3).unwrap()), + vec![7, 7, 7] + ); +} + +#[test] +fn metadata_only_rle_round_trip() { + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::constant(Some(10_u64.to_le_bytes().to_vec().into())), + ProtobufUtils21::constant(Some(2_u32.to_le_bytes().to_vec().into())), + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(!has_payload); + assert_eq!( + decoded_u64(decoder.decompress(None, 6).unwrap()), + vec![10; 6] + ); + assert!(decoder.decompress(None, 5).is_err()); +} + +#[test] +fn rle_framing_is_fallible() { + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(64, None), + ProtobufUtils21::constant(Some(vec![2_u8].into())), + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(has_payload); + assert!( + decoder + .decompress(Some(LanceBuffer::from(vec![0; 7])), 4) + .is_err() + ); + + let mut payload = 100_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&[0; 8]); + assert!( + decoder + .decompress(Some(LanceBuffer::from(payload)), 4) + .is_err() + ); + + let mut payload = 16_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&[0; 16]); + payload.push(1); + let error = decoder + .decompress(Some(LanceBuffer::from(payload)), 4) + .unwrap_err(); + assert!(error.to_string().contains("Metadata-only RLE run-length")); +} + +#[test] +fn decoder_allocation_overflow_is_fallible() { + let (decoder, has_payload) = create_block_decompressor( + &ProtobufUtils21::constant(Some(1_u64.to_le_bytes().to_vec().into())), + BlockValueType::UInt64, + ) + .unwrap(); + assert!(!has_payload); + assert!(decoder.decompress(None, u64::MAX).is_err()); +} + +#[cfg(feature = "bitpacking")] +#[test] +fn frozen_out_of_line_compressor_rejects_wider_reuse() { + use crate::encodings::physical::bitpacking::OutOfLineBitpacking; + + let compressor = OutOfLineBitpacking::new(3, 64); + let error = compressor + .compress(DataBlock::FixedWidth(fixed_u64(&[1, 2, 8]))) + .unwrap_err(); + assert!(error.to_string().contains("requires 4 bits")); +} diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index cdcb0e0c482..cad9ade7d54 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -45,7 +45,7 @@ use crate::utils::bytepack::ByteUnpacker; use crate::{ compression::{ BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor, - create_rle_decompressor, + compress_required_block, create_rle_decompressor, }, data::{AllNullDataBlock, DataBlock, VariableWidthBlock}, utils::bytepack::BytepackedIntegerEncoder, @@ -179,7 +179,7 @@ impl DecodeMiniBlockTask { levels: LanceBuffer, num_levels: u16, ) -> Result> { - let rep = rep_decompressor.decompress(levels, num_levels as u64)?; + let rep = rep_decompressor.decompress(Some(levels), num_levels as u64)?; let rep = rep.as_fixed_width().unwrap(); debug_assert_eq!(rep.num_values, num_levels as u64); debug_assert_eq!(rep.bits_per_value, 16); @@ -1556,7 +1556,7 @@ impl StructuralPageScheduler for ComplexAllNullScheduler { } LevelCodec::Block(decompressor) => { let frame = LanceBuffer::from_bytes(compressed_bytes, 1); - let decompressed = decompressor.decompress(frame, num_values)?; + let decompressed = decompressor.decompress(Some(frame), num_values)?; dense_levels_from_block(decompressed, num_values, level_type) } } @@ -2580,7 +2580,10 @@ impl StructuralPageScheduler for MiniBlockScheduler { let dictionary = if let Some(ref mut dictionary) = self.dictionary { let dictionary_data = dictionary_bytes.unwrap(); Some(Arc::new(dictionary.dictionary_decompressor.decompress( - LanceBuffer::from_bytes(dictionary_data, dictionary.dictionary_data_alignment), + Some(LanceBuffer::from_bytes( + dictionary_data, + dictionary.dictionary_data_alignment, + )), dictionary.num_dictionary_items, )?)) } else { @@ -4899,7 +4902,11 @@ impl PrimitiveStructuralEncoder { }; chunk_fixed_width.compute_stat(); let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width); - let compressed_levels = compressor.compress(chunk_levels_block)?; + let compressed_levels = compressor.compress(chunk_levels_block)?.ok_or_else(|| { + Error::internal( + "Rep/def block compressor selected a metadata-only codec".to_string(), + ) + })?; let num_levels = u16::try_from(num_chunk_levels).map_err(|_| { Error::invalid_input_source( format!( @@ -4960,9 +4967,8 @@ impl PrimitiveStructuralEncoder { let levels_block = DataBlock::FixedWidth(fixed_width_block); let levels_field = Field::new_arrow("", DataType::UInt16, false)?; - let (compressor, encoding) = - compression_strategy.create_block_compressor(&levels_field, &levels_block)?; - let compressed_buffer = compressor.compress(levels_block)?; + let (compressed_buffer, encoding) = + compress_required_block(compression_strategy, &levels_field, levels_block)?; Ok((compressed_buffer, encoding)) } @@ -5346,9 +5352,8 @@ impl PrimitiveStructuralEncoder { let num_dictionary_items = dictionary_data.num_values(); let dict_values_field = Self::build_dict_values_compressor_field(field)?; - let (compressor, dictionary_encoding) = compression_strategy - .create_block_compressor(&dict_values_field, &dictionary_data)?; - let dictionary_buffer = compressor.compress(dictionary_data)?; + let (dictionary_buffer, dictionary_encoding) = + compress_required_block(compression_strategy, &dict_values_field, dictionary_data)?; data.push(dictionary_buffer); if let Some(rep_index) = rep_index { @@ -8914,7 +8919,7 @@ mod tests { .create_block_decompressor(&encoding) .unwrap(); let decompressed = decompressor - .decompress(compressed_buf, values.len() as u64) + .decompress(Some(compressed_buf), values.len() as u64) .unwrap(); let decompressed_fixed_width = decompressed.as_fixed_width().unwrap(); assert_eq!(decompressed_fixed_width.num_values, values.len() as u64); @@ -9002,6 +9007,7 @@ mod tests { }); BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block) .unwrap() + .unwrap() } fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns { diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs index 68becc0014f..c0f6901d962 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs @@ -1069,7 +1069,6 @@ enum SparsePositionSetDecoder { }, Explicit { decompressor: Arc, - encoding: CompressiveEncoding, count: u64, domain_len: u64, }, @@ -1084,7 +1083,6 @@ enum SparseCountSetDecoder { }, Explicit { decompressor: Arc, - encoding: CompressiveEncoding, count: u64, }, } @@ -1297,7 +1295,7 @@ impl SparseStructuralScheduler { let layer_decompressors = layout .structural_layers .iter() - .map(|layer| Self::layer_decompressors(layer, decompressors)) + .map(Self::layer_decompressors) .collect::>>()?; Ok(Self { @@ -2095,35 +2093,30 @@ impl SparseStructuralScheduler { fn create_position_decompressor( compression: &CompressiveEncoding, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result> { let compression = Self::validate_compression(compression, label)?; - Self::validate_block_encoding(compression, label)?; - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - decompressors.create_block_decompressor(compression) - })) - .map_err(|_| { - Error::invalid_input_source( - format!( - "Sparse structural {label} descriptor caused decompressor construction to panic" - ) - .into(), - ) - })? - .map(Arc::from) + let (decompressor, has_payload) = crate::compression::block::create_block_decompressor( + compression, + crate::compression::BlockValueType::UInt64, + ) .map_err(|error| { Error::invalid_input_source( - format!("Sparse structural {label} decompressor construction failed: {error}") - .into(), + format!("Sparse structural {label} block validation failed: {error}").into(), ) - }) + })?; + if !has_payload { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} explicit encoding must require one payload") + .into(), + )); + } + Ok(Arc::from(decompressor)) } fn position_set_decoder( position_set: &pb21::SparsePositionSet, domain_len: u64, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result<(SparsePositionSetDecoder, u64)> { let cardinality = Self::position_cardinality(position_set, domain_len, label)?; let positions = position_set.positions.as_ref().ok_or_else(|| { @@ -2145,12 +2138,7 @@ impl SparseStructuralScheduler { } pb21::sparse_position_set::Positions::Explicit(compression) => { SparsePositionSetDecoder::Explicit { - decompressor: Self::create_position_decompressor( - compression, - label, - decompressors, - )?, - encoding: compression.clone(), + decompressor: Self::create_position_decompressor(compression, label)?, count: cardinality, domain_len, } @@ -2164,7 +2152,6 @@ impl SparseStructuralScheduler { validity_set: &pb21::SparseValiditySet, domain_len: u64, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result<(SparseValiditySetDecoder, u64)> { let meaning = Self::validity_meaning(validity_set, label)?; let position_set = validity_set.positions.as_ref().ok_or_else(|| { @@ -2172,8 +2159,7 @@ impl SparseStructuralScheduler { format!("Sparse structural {label} positions are required").into(), ) })?; - let (positions, cardinality) = - Self::position_set_decoder(position_set, domain_len, label, decompressors)?; + let (positions, cardinality) = Self::position_set_decoder(position_set, domain_len, label)?; Ok((SparseValiditySetDecoder { meaning, positions }, cardinality)) } @@ -2203,7 +2189,6 @@ impl SparseStructuralScheduler { count_set: &pb21::SparseCountSet, cardinality: u64, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result { Self::count_buffer_count(count_set, cardinality, label)?; let counts = count_set.counts.as_ref().ok_or_else(|| { @@ -2219,12 +2204,7 @@ impl SparseStructuralScheduler { }, pb21::sparse_count_set::Counts::Explicit(compression) => { SparseCountSetDecoder::Explicit { - decompressor: Self::create_position_decompressor( - compression, - label, - decompressors, - )?, - encoding: compression.clone(), + decompressor: Self::create_position_decompressor(compression, label)?, count: cardinality, } } @@ -2335,7 +2315,6 @@ impl SparseStructuralScheduler { fn layer_decompressors( layer: &pb21::SparseStructuralLayer, - decompressors: &dyn DecompressionStrategy, ) -> Result { Ok(match Self::require_layer(layer)? { pb21::sparse_structural_layer::Layer::Validity(layer) => { @@ -2343,7 +2322,6 @@ impl SparseStructuralScheduler { Self::require_validity_set(&layer.validity, "validity")?, layer.num_slots, "validity positions", - decompressors, )?; SparseLayerDecompressors::Validity { num_slots: layer.num_slots, @@ -2357,19 +2335,16 @@ impl SparseStructuralScheduler { non_empty_positions, layer.num_slots, "list non-empty positions", - decompressors, )?; let counts = Self::count_set_decoder( Self::require_count_set(&layer.counts, "list counts")?, num_non_empty, "list counts", - decompressors, )?; let (validity, _) = Self::validity_set_decoder( Self::require_validity_set(&layer.validity, "list")?, layer.num_slots, "list validity positions", - decompressors, )?; SparseLayerDecompressors::List { num_slots: layer.num_slots, @@ -2384,7 +2359,6 @@ impl SparseStructuralScheduler { Self::require_validity_set(&layer.validity, "fixed-size-list")?, layer.num_slots, "fixed-size-list validity positions", - decompressors, )?; SparseLayerDecompressors::FixedSizeList { num_slots: layer.num_slots, @@ -2575,88 +2549,19 @@ impl SparseStructuralScheduler { Ok(()) } - fn validate_structural_buffer_headers( - encoding: &CompressiveEncoding, - data: &[u8], - label: &str, - ) -> Result<()> { - use pb21::compressive_encoding::Compression; - - match encoding.compression.as_ref() { - Some(Compression::General(general)) => { - Self::validate_general_buffer_header(general, data, label) - } - Some(Compression::Rle(rle)) => { - let values_size = u64::from_le_bytes( - data.get(..8) - .ok_or_else(|| { - Error::invalid_input_source( - format!( - "Sparse structural {label} RLE buffer is missing its header" - ) - .into(), - ) - })? - .try_into() - .map_err(|_| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE header is malformed").into(), - ) - })?, - ); - let values_size = usize_from_u64(values_size, "RLE values buffer size")?; - let values_end = 8_usize.checked_add(values_size).ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE values range overflows").into(), - ) - })?; - let values_data = data.get(8..values_end).ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE values buffer is truncated").into(), - ) - })?; - let lengths_data = data.get(values_end..).ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE run-length buffer is missing") - .into(), - ) - })?; - Self::validate_general_child_buffer( - Self::require_encoding(&rle.values, "RLE values")?, - values_data, - "RLE values", - )?; - Self::validate_general_child_buffer( - Self::require_encoding(&rle.run_lengths, "RLE run lengths")?, - lengths_data, - "RLE run lengths", - ) - } - _ => Ok(()), - } - } - fn decode_u64_values( decompressor: &dyn BlockDecompressor, - encoding: &CompressiveEncoding, data: Bytes, num_values: u64, label: &str, ) -> Result> { - Self::validate_structural_buffer_headers(encoding, &data, label)?; - let decoded = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - decompressor.decompress(LanceBuffer::from_bytes(data, 1), num_values) - })) - .map_err(|_| { - Error::invalid_input_source( - format!("Sparse structural {label} decompression panicked").into(), - ) - })? - .map_err(|error| { - Error::invalid_input_source( - format!("Sparse structural {label} decompression failed: {error}").into(), - ) - })?; + let decoded = decompressor + .decompress(Some(LanceBuffer::from_bytes(data, 1)), num_values) + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} decompression failed: {error}").into(), + ) + })?; let fixed = decoded.as_fixed_width().ok_or_else(|| { Error::invalid_input_source( format!("Sparse structural {label} did not decode to fixed width data").into(), @@ -2718,14 +2623,12 @@ impl SparseStructuralScheduler { fn decode_explicit_positions( decompressor: &Arc, - encoding: &CompressiveEncoding, data: Bytes, num_positions: u64, num_slots: u64, label: &str, ) -> Result { - let deltas = - Self::decode_u64_values(decompressor.as_ref(), encoding, data, num_positions, label)?; + let deltas = Self::decode_u64_values(decompressor.as_ref(), data, num_positions, label)?; let mut positions = Vec::with_capacity(deltas.len()); let mut current = 0_u64; for (idx, delta) in deltas.into_iter().enumerate() { @@ -2771,12 +2674,10 @@ impl SparseStructuralScheduler { } SparsePositionSetDecoder::Explicit { decompressor, - encoding, count, domain_len, } => Self::decode_explicit_positions( decompressor, - encoding, Self::next_structural_buffer(buffers, label)?, *count, *domain_len, @@ -2808,12 +2709,10 @@ impl SparseStructuralScheduler { } SparseCountSetDecoder::Explicit { decompressor, - encoding, count, } => { let counts = Self::decode_u64_values( decompressor.as_ref(), - encoding, Self::next_structural_buffer(buffers, label)?, *count, label, @@ -4684,11 +4583,9 @@ fn slice_sparse_plan( mod tests { use std::sync::Mutex; - use crate::{ - compression::DefaultDecompressionStrategy, - encodings::physical::block::{CompressionConfig, CompressionScheme}, - testing::SimulatedScheduler, - }; + #[cfg(feature = "lz4")] + use crate::encodings::physical::block::{CompressionConfig, CompressionScheme}; + use crate::{compression::DefaultDecompressionStrategy, testing::SimulatedScheduler}; use super::*; @@ -4723,6 +4620,7 @@ mod tests { ) } + #[cfg(feature = "lz4")] fn general_lz4(values: CompressiveEncoding) -> CompressiveEncoding { ProtobufUtils21::wrapped(CompressionConfig::new(CompressionScheme::Lz4, None), values) .unwrap() diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index 809536a2c1b..fc3f3416cba 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -10,7 +10,7 @@ use lance_core::{Error, Result, datatypes::Field, utils::bit::pad_bytes}; use crate::{ buffer::LanceBuffer, - compression::CompressionStrategy, + compression::{CompressionStrategy, compress_required_block}, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, decoder::PageEncoding, encoder::EncodedPage, @@ -588,8 +588,7 @@ fn encode_u64_values( }); block.compute_stat(); let field = Field::new_arrow("", arrow_schema::DataType::UInt64, false)?; - let (compressor, encoding) = compression_strategy.create_block_compressor(&field, &block)?; - Ok((compressor.compress(block)?, encoding)) + compress_required_block(compression_strategy, &field, block) } fn positions_to_deltas(positions: &[u64], label: &str) -> Result> { @@ -887,8 +886,9 @@ mod tests { use crate::{ constants::{ - PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, - STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, + COMPRESSION_META_KEY, PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, + STRUCTURAL_ENCODING_SPARSE, }, data::FixedSizeListBlock, encoder::{ @@ -1545,6 +1545,62 @@ mod tests { check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; } + #[tokio::test] + async fn test_v2_3_sparse_variable_values_keep_legacy_flat_offsets() { + let num_values = 4_096; + let values = Arc::new(StringArray::from_iter_values((0..num_values).map( + |index| { + let length = [4_usize, 10, 5, 8][index % 4]; + format!("{index:04x}{}", "x".repeat(length - 4)) + }, + ))) as ArrayRef; + let array = + sparse_list_values( + num_values * 2, + 2, + values, + Arc::new(ArrowField::new("item", DataType::Utf8, true).with_metadata( + HashMap::from([(COMPRESSION_META_KEY.to_string(), "none".to_string())]), + )), + ); + let metadata = sparse_metadata(); + let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, metadata.clone()) + .await + .unwrap(); + assert_eq!(pages.len(), 1); + let sparse = sparse_layout(&pages[0]); + assert_eq!(sparse.num_buffers, 1); + let Some(pb21::compressive_encoding::Compression::Variable(variable)) = sparse + .value_compression + .as_ref() + .and_then(|encoding| encoding.compression.as_ref()) + else { + panic!( + "expected sparse variable value compression, got {:?}", + sparse.value_compression + ); + }; + assert!(variable.values.is_none()); + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(pb21::compressive_encoding::Compression::Flat(pb21::Flat { + bits_per_value: 32, + data: None, + })) + )); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..17) + .with_indices(vec![0, 7, (num_values * 2 - 1) as u64]); + check_round_trip_encoding_of_data(vec![array], &cases, metadata).await; + } + #[tokio::test] async fn test_explicit_sparse_struct_with_constant_and_sparse_children() { let fields = Fields::from(vec![ diff --git a/rust/lance-encoding/src/encodings/physical.rs b/rust/lance-encoding/src/encodings/physical.rs index 0439c0216fb..ea13dad7970 100644 --- a/rust/lance-encoding/src/encodings/physical.rs +++ b/rust/lance-encoding/src/encodings/physical.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::{Error, Result}; + pub mod binary; #[cfg(feature = "bitpacking")] pub mod bitpacking; @@ -12,3 +14,33 @@ pub mod general; pub mod packed; pub mod rle; pub mod value; + +pub(crate) fn checked_vec_capacity( + num_values: u64, + bytes_per_value: usize, + label: &str, +) -> Result { + let capacity = usize::try_from(num_values) + .map_err(|_| Error::invalid_input(format!("{label} cardinality does not fit usize")))?; + let output_bytes = capacity + .checked_mul(bytes_per_value) + .ok_or_else(|| Error::invalid_input(format!("{label} byte length overflows usize")))?; + if output_bytes > isize::MAX as usize { + return Err(Error::invalid_input(format!( + "{label} byte length {output_bytes} exceeds isize::MAX" + ))); + } + Ok(capacity) +} + +pub(crate) fn try_vec_with_capacity(num_values: u64, label: &str) -> Result> { + let capacity = checked_vec_capacity(num_values, std::mem::size_of::(), label)?; + let output_bytes = capacity * std::mem::size_of::(); + let mut values = Vec::new(); + values.try_reserve_exact(capacity).map_err(|error| { + Error::invalid_input(format!( + "{label} could not reserve {capacity} values ({output_bytes} bytes): {error}" + )) + })?; + Ok(values) +} diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index 7c84c2774b7..450eff49195 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -15,6 +15,7 @@ use core::panic; use crate::compression::{ BlockCompressor, BlockDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor, + require_block_payload, }; use crate::buffer::LanceBuffer; @@ -275,20 +276,35 @@ impl BinaryMiniBlockDecompressor { Self { bits_per_offset } } - pub fn from_variable(variable: &pb21::Variable) -> Self { - if let Compression::Flat(flat) = variable + pub fn from_variable(variable: &pb21::Variable) -> Result { + if variable.values.is_some() { + return Err(Error::invalid_input( + "Binary mini-block Variable values encoding must be absent", + )); + } + let offsets = variable .offsets .as_ref() - .unwrap() + .ok_or_else(|| Error::invalid_input("Variable encoding is missing offsets"))?; + let compression = offsets .compression .as_ref() - .unwrap() - { - Self { - bits_per_offset: flat.bits_per_value as u8, + .ok_or_else(|| Error::invalid_input("Variable offsets are missing compression"))?; + match compression { + Compression::Flat(flat) + if matches!(flat.bits_per_value, 32 | 64) && flat.data.is_none() => + { + Ok(Self { + bits_per_offset: flat.bits_per_value as u8, + }) } - } else { - panic!("Unsupported offsets compression: {:?}", variable.offsets); + Compression::Flat(flat) => Err(Error::invalid_input(format!( + "Variable offsets require uncompressed 32 or 64-bit Flat encoding, got {} bits", + flat.bits_per_value + ))), + other => Err(Error::invalid_input(format!( + "Unsupported legacy variable offset compression: {other:?}" + ))), } } } @@ -361,7 +377,7 @@ impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { pub struct VariableEncoder {} impl BlockCompressor for VariableEncoder { - fn compress(&self, mut data: DataBlock) -> Result { + fn compress(&self, mut data: DataBlock) -> Result> { match data { DataBlock::VariableWidth(ref mut variable_width_data) => { match variable_width_data.bits_per_offset { @@ -413,18 +429,17 @@ impl BlockCompressor for VariableEncoder { output.extend_from_slice(&variable_width_data.data); Ok(LanceBuffer::from(output)) } - _ => { - panic!( - "BinaryBlockEncoder does not work with {} bits per offset VariableWidth DataBlock.", - variable_width_data.bits_per_offset - ); - } + _ => Err(Error::invalid_input(format!( + "BinaryBlockEncoder does not support {}-bit offsets", + variable_width_data.bits_per_offset + ))), } } - _ => { - panic!("BinaryBlockEncoder can only work with Variable Width DataBlock."); - } + _ => Err(Error::invalid_input( + "BinaryBlockEncoder requires a variable-width block", + )), } + .map(Some) } } @@ -455,7 +470,8 @@ impl VariablePerValueDecompressor for VariableDecoder { pub struct BinaryBlockDecompressor {} impl BlockDecompressor for BinaryBlockDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Binary block")?; // In older (not quite stable) versions we stored the bits per offset as a single byte and then the num_values // as four bytes. However, this led to alignment problems and was wasteful since we already store the num_values // in higher layers. diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index ad1b567edd7..9f049758c80 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -16,7 +16,7 @@ //! we can easily jump to the correct value. use arrow_array::types::UInt64Type; -use arrow_array::{Array, PrimitiveArray}; +use arrow_array::{Array, PrimitiveArray, UInt64Array}; use arrow_buffer::ArrowNativeType; use byteorder::{ByteOrder, LittleEndian}; use lance_bitpacking::BitPacking; @@ -24,7 +24,13 @@ use lance_bitpacking::BitPacking; use lance_core::{Error, Result}; use crate::buffer::LanceBuffer; -use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor}; +use crate::compression::{ + BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, + block::{ + validate_inline_bitpacking_payload, validate_out_of_line_payload, visit_unsigned_values, + }, + require_block_payload, +}; use crate::data::BlockInfo; use crate::data::{DataBlock, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ @@ -34,10 +40,17 @@ use crate::format::pb21::CompressiveEncoding; use crate::format::{ProtobufUtils21, pb21}; use crate::statistics::{GetStat, Stat}; use bytemuck::{AnyBitPattern, cast_slice}; +#[cfg(test)] +use std::cell::Cell; const LOG_ELEMS_PER_CHUNK: u8 = 10; const ELEMS_PER_CHUNK: u64 = 1 << LOG_ELEMS_PER_CHUNK; +#[cfg(test)] +thread_local! { + pub(crate) static BLOCK_WIDTH_VALIDATION_COUNT: Cell = const { Cell::new(0) }; +} + #[derive(Debug, Default)] pub struct InlineBitpacking { uncompressed_bit_width: u64, @@ -234,10 +247,42 @@ impl MiniBlockCompressor for InlineBitpacking { } impl BlockCompressor for InlineBitpacking { - fn compress(&self, data: DataBlock) -> Result { - let fixed_width = data.as_fixed_width().unwrap(); + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(fixed_width) = data else { + return Err(Error::invalid_input( + "Inline bitpacking requires fixed-width data", + )); + }; + if fixed_width.bits_per_value != self.uncompressed_bit_width { + return Err(Error::invalid_input(format!( + "Inline bitpacking expects {}-bit values, got {}", + self.uncompressed_bit_width, fixed_width.bits_per_value + ))); + } + let value_type = BlockValueType::from_bits(self.uncompressed_bit_width)?; + if fixed_width.num_values == 0 || fixed_width.num_values > ELEMS_PER_CHUNK { + return Err(Error::invalid_input(format!( + "Inline block bitpacking requires 1..={ELEMS_PER_CHUNK} values, got {}", + fixed_width.num_values + ))); + } + let mut max_value = 0_u64; + visit_unsigned_values(&fixed_width, value_type, |value| { + max_value = max_value.max(value); + Ok(()) + })?; + let compressed_bit_width = u64::from(u64::BITS - max_value.leading_zeros()).max(1); + fixed_width.block_info.0.write().unwrap().insert( + Stat::BitWidth, + std::sync::Arc::new(UInt64Array::from(vec![compressed_bit_width])), + ); let (chunked, _) = self.chunk_data(fixed_width); - Ok(chunked.data.into_iter().next().unwrap()) + chunked + .data + .into_iter() + .next() + .map(Some) + .ok_or_else(|| Error::internal("Inline bitpacking produced no payload".to_string())) } } @@ -265,7 +310,13 @@ impl MiniBlockDecompressor for InlineBitpacking { } impl BlockDecompressor for InlineBitpacking { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Inline bitpacking")?; + validate_inline_bitpacking_payload( + &data, + BlockValueType::from_bits(self.uncompressed_bit_width)?, + num_values, + )?; match self.uncompressed_bit_width { 8 => Self::unchunk::(data, num_values), 16 => Self::unchunk::(data, num_values), @@ -477,27 +528,45 @@ impl OutOfLineBitpacking { } impl BlockCompressor for OutOfLineBitpacking { - fn compress(&self, data: DataBlock) -> Result { - let fixed_width = data.as_fixed_width().unwrap(); + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(fixed_width) = data else { + return Err(Error::invalid_input( + "Out-of-line bitpacking requires fixed-width data", + )); + }; + if fixed_width.bits_per_value != self.uncompressed_bit_width { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking expects {}-bit values, got {}", + self.uncompressed_bit_width, fixed_width.bits_per_value + ))); + } + validate_frozen_block_width(&fixed_width, self.compressed_bit_width)?; let compressed = match fixed_width.bits_per_value { 8 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), 16 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), 32 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), 64 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), - _ => panic!("Bitpacking word size must be 8,16,32,64"), + _ => unreachable!("bitpacking word size was validated"), }; - Ok(compressed) + Ok(Some(compressed)) } } impl BlockDecompressor for OutOfLineBitpacking { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Out-of-line bitpacking")?; + validate_out_of_line_payload( + &data, + BlockValueType::from_bits(self.uncompressed_bit_width)?, + num_values, + self.compressed_bit_width, + )?; let word_size = match self.uncompressed_bit_width { 8 => std::mem::size_of::(), 16 => std::mem::size_of::(), 32 => std::mem::size_of::(), 64 => std::mem::size_of::(), - _ => panic!("Bitpacking word size must be 8,16,32,64"), + _ => unreachable!("bitpacking word size was validated"), }; debug_assert_eq!(data.len() % word_size, 0); let total_words = (data.len() / word_size) as u64; @@ -535,6 +604,38 @@ impl BlockDecompressor for OutOfLineBitpacking { } } +fn validate_frozen_block_width( + block: &FixedWidthDataBlock, + compressed_bits_per_value: u64, +) -> Result<()> { + #[cfg(test)] + BLOCK_WIDTH_VALIDATION_COUNT.with(|count| count.set(count.get().saturating_add(1))); + + let value_type = BlockValueType::from_bits(block.bits_per_value)?; + if compressed_bits_per_value >= value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking width {compressed_bits_per_value} must be less than the {}-bit input width", + value_type.bits_per_value() + ))); + } + let max_value = if compressed_bits_per_value == 0 { + 0 + } else { + (1_u64 << compressed_bits_per_value) - 1 + }; + let mut index = 0_u64; + visit_unsigned_values(block, value_type, |value| { + if value > max_value { + let required_bits = u64::from(u64::BITS - value.leading_zeros()); + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking codec uses {compressed_bits_per_value} bits but input value {value} at index {index} requires {required_bits} bits" + ))); + } + index += 1; + Ok(()) + }) +} + #[cfg(test)] mod test { use std::{collections::HashMap, sync::Arc}; @@ -546,7 +647,7 @@ mod test { use super::{ELEMS_PER_CHUNK, InlineBitpacking, bitpack_out_of_line, unpack_out_of_line}; use crate::{ buffer::LanceBuffer, - compression::MiniBlockDecompressor, + compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor}, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, testing::{TestCases, check_round_trip_encoding_of_data}, version::LanceFileVersion, @@ -571,6 +672,27 @@ mod test { assert_eq!(block.data.len(), 0); } + #[test] + fn test_inline_block_bitpacking_computes_its_own_bit_width() { + let codec = InlineBitpacking::new(32); + let input = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_u32, 1, 7]), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::new(), + }); + + let payload = BlockCompressor::compress(&codec, input).unwrap().unwrap(); + let decoded = BlockDecompressor::decompress(&codec, Some(payload), 3).unwrap(); + let DataBlock::FixedWidth(decoded) = decoded else { + panic!("Expected FixedWidth block"); + }; + assert_eq!( + decoded.data.borrow_to_typed_view::().as_ref(), + &[0, 1, 7] + ); + } + #[test_log::test(tokio::test)] async fn test_miniblock_bitpack() { let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); diff --git a/rust/lance-encoding/src/encodings/physical/block.rs b/rust/lance-encoding/src/encodings/physical/block.rs index a1f5bdb3fdd..17af022ae71 100644 --- a/rust/lance-encoding/src/encodings/physical/block.rs +++ b/rust/lance-encoding/src/encodings/physical/block.rs @@ -34,7 +34,7 @@ use crate::format::{ }; use crate::{ buffer::LanceBuffer, - compression::VariablePerValueDecompressor, + compression::{VariablePerValueDecompressor, require_block_payload}, data::{BlockInfo, DataBlock, VariableWidthBlock}, encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}, }; @@ -130,12 +130,40 @@ impl FromStr for CompressionScheme { pub trait BufferCompressor: std::fmt::Debug + Send + Sync { fn compress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()>; fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()>; + /// Decompresses at most `max_output_len` bytes. + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()>; + /// Decompresses exactly `expected_output_len` bytes without trusting an + /// embedded size prefix to choose the allocation size. + fn decompress_exact( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + expected_output_len: usize, + ) -> Result<()> { + let start = output_buf.len(); + self.decompress_bounded(input_buf, output_buf, expected_output_len)?; + let actual_output_len = output_buf.len().checked_sub(start).ok_or_else(|| { + Error::internal("Buffer decompressor shortened its output".to_string()) + })?; + if actual_output_len != expected_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Decompression produced {actual_output_len} bytes, expected exactly {expected_output_len}" + ))); + } + Ok(()) + } fn config(&self) -> CompressionConfig; } #[cfg(feature = "zstd")] mod zstd { - use std::io::{Cursor, Write}; + use std::io::{Cursor, Read, Write}; use std::sync::{Mutex, OnceLock}; use super::*; @@ -279,6 +307,101 @@ mod zstd { Ok(()) } + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()> { + if self.is_raw_stream_format(input_buf) { + let read_limit = max_output_len.checked_add(1).ok_or_else(|| { + Error::invalid_input("Zstd maximum output length overflows usize") + })?; + let read_limit_u64 = u64::try_from(read_limit).map_err(|_| { + Error::invalid_input("Zstd maximum output length does not fit u64") + })?; + let start = output_buf.len(); + output_buf.try_reserve_exact(read_limit).map_err(|error| { + Error::invalid_input(format!( + "Zstd could not reserve {read_limit} bounded output bytes: {error}" + )) + })?; + let result = (|| { + let decoder = ::zstd::stream::read::Decoder::new(Cursor::new(input_buf))?; + let mut bounded = decoder.take(read_limit_u64); + bounded.read_to_end(output_buf) + })(); + let actual_output_len = match result { + Ok(actual_output_len) => actual_output_len, + Err(error) => { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd decompression failed: {error}" + ))); + } + }; + if actual_output_len > max_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd output exceeds the {max_output_len}-byte limit" + ))); + } + return Ok(()); + } + + const LENGTH_PREFIX_SIZE: usize = 8; + if input_buf.len() < LENGTH_PREFIX_SIZE { + return Err(Error::invalid_input(format!( + "Length-prefixed Zstd payload has {} bytes, shorter than its {LENGTH_PREFIX_SIZE}-byte prefix", + input_buf.len() + ))); + } + let declared_output_len = + u64::from_le_bytes(input_buf[..LENGTH_PREFIX_SIZE].try_into().map_err(|_| { + Error::invalid_input("Length-prefixed Zstd size prefix is truncated") + })?); + let declared_output_len = usize::try_from(declared_output_len).map_err(|_| { + Error::invalid_input("Length-prefixed Zstd output length does not fit usize") + })?; + if declared_output_len > max_output_len { + return Err(Error::invalid_input(format!( + "Length-prefixed Zstd payload declares {declared_output_len} output bytes, exceeding the {max_output_len}-byte limit" + ))); + } + + let start = output_buf.len(); + let end = start + .checked_add(declared_output_len) + .ok_or_else(|| Error::invalid_input("Zstd output buffer length overflows usize"))?; + output_buf + .try_reserve_exact(declared_output_len) + .map_err(|error| { + Error::invalid_input(format!( + "Zstd could not reserve {declared_output_len} output bytes: {error}" + )) + })?; + output_buf.resize(end, 0); + let decompressed_size = match decompress_to_buffer( + &input_buf[LENGTH_PREFIX_SIZE..], + &mut output_buf[start..end], + ) { + Ok(decompressed_size) => decompressed_size, + Err(error) => { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd decompression failed: {error}" + ))); + } + }; + if decompressed_size != declared_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd decompressed {decompressed_size} bytes, but its prefix declares {declared_output_len}" + ))); + } + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::Zstd, @@ -347,6 +470,62 @@ mod lz4 { Ok(()) } + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()> { + const LENGTH_PREFIX_SIZE: usize = 4; + if input_buf.len() < LENGTH_PREFIX_SIZE { + return Err(Error::invalid_input(format!( + "LZ4 payload has {} bytes, shorter than its {LENGTH_PREFIX_SIZE}-byte prefix", + input_buf.len() + ))); + } + let declared_output_len = u32::from_le_bytes( + input_buf[..LENGTH_PREFIX_SIZE] + .try_into() + .map_err(|_| Error::invalid_input("LZ4 size prefix is truncated"))?, + ) as usize; + if declared_output_len > max_output_len { + return Err(Error::invalid_input(format!( + "LZ4 payload declares {declared_output_len} output bytes, exceeding the {max_output_len}-byte limit" + ))); + } + + let start = output_buf.len(); + let end = start + .checked_add(declared_output_len) + .ok_or_else(|| Error::invalid_input("LZ4 output buffer length overflows usize"))?; + output_buf + .try_reserve_exact(declared_output_len) + .map_err(|error| { + Error::invalid_input(format!( + "LZ4 could not reserve {declared_output_len} output bytes: {error}" + )) + })?; + output_buf.resize(end, 0); + let decompressed_size = + match ::lz4::block::decompress_to_buffer(input_buf, None, &mut output_buf[start..]) + { + Ok(decompressed_size) => decompressed_size, + Err(error) => { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "LZ4 decompression failed: {error}" + ))); + } + }; + if decompressed_size != declared_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "LZ4 decompressed {decompressed_size} bytes, but its prefix declares {declared_output_len}" + ))); + } + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::Lz4, @@ -370,6 +549,30 @@ impl BufferCompressor for NoopBufferCompressor { Ok(()) } + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()> { + if input_buf.len() > max_output_len { + return Err(Error::invalid_input(format!( + "Uncompressed payload has {} bytes, exceeding the {max_output_len}-byte limit", + input_buf.len() + ))); + } + output_buf + .try_reserve_exact(input_buf.len()) + .map_err(|error| { + Error::invalid_input(format!( + "Uncompressed payload could not reserve {} output bytes: {error}", + input_buf.len() + )) + })?; + output_buf.extend_from_slice(input_buf); + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::None, @@ -439,11 +642,12 @@ impl GeneralBlockDecompressor { } impl BlockDecompressor for GeneralBlockDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "General block compression")?; let mut decompressed = Vec::new(); self.compressor.decompress(&data, &mut decompressed)?; self.inner - .decompress(LanceBuffer::from(decompressed), num_values) + .decompress(Some(LanceBuffer::from(decompressed)), num_values) } } @@ -603,13 +807,19 @@ impl VariablePerValueDecompressor for CompressedBufferEncoder { } impl BlockCompressor for CompressedBufferEncoder { - fn compress(&self, data: DataBlock) -> Result { + fn compress(&self, data: DataBlock) -> Result> { let encoded = match data { DataBlock::FixedWidth(fixed_width) => fixed_width.data, DataBlock::VariableWidth(variable_width) => { // Wrap VariableEncoder to handle the encoding let encoder = VariableEncoder::default(); BlockCompressor::compress(&encoder, DataBlock::VariableWidth(variable_width))? + .ok_or_else(|| { + Error::internal( + "VariableEncoder returned no payload for general compression" + .to_string(), + ) + })? } _ => { return Err(Error::invalid_input_source( @@ -620,18 +830,19 @@ impl BlockCompressor for CompressedBufferEncoder { let mut compressed = Vec::new(); self.compressor.compress(&encoded, &mut compressed)?; - Ok(LanceBuffer::from(compressed)) + Ok(Some(LanceBuffer::from(compressed))) } } impl BlockDecompressor for CompressedBufferEncoder { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Compressed variable block")?; let mut decompressed = Vec::new(); self.compressor.decompress(&data, &mut decompressed)?; // Delegate to BinaryBlockDecompressor which handles the inline metadata let inner_decoder = BinaryBlockDecompressor::default(); - inner_decoder.decompress(LanceBuffer::from(decompressed), num_values) + inner_decoder.decompress(Some(LanceBuffer::from(decompressed)), num_values) } } @@ -640,6 +851,7 @@ mod tests { use super::*; use std::str::FromStr; + #[cfg(feature = "zstd")] use crate::encodings::physical::block::zstd::ZstdBufferCompressor; #[test] @@ -679,6 +891,32 @@ mod tests { .decompress(&compressed_data, &mut decompressed_data) .unwrap(); assert_eq!(input_data, decompressed_data.as_slice()); + + let mut exact = Vec::new(); + compressor + .decompress_exact(&compressed_data, &mut exact, input_data.len()) + .unwrap(); + assert_eq!(input_data, exact.as_slice()); + } + + #[test] + fn test_length_prefixed_zstd_rejects_untrusted_output_size() { + let compressor = ZstdBufferCompressor::new(0); + let mut compressed_data = Vec::new(); + compressor + .compress(b"bounded", &mut compressed_data) + .unwrap(); + compressed_data[..8].copy_from_slice(&u64::MAX.to_le_bytes()); + + let mut output = vec![7, 9]; + let error = compressor + .decompress_exact(&compressed_data, &mut output, 7) + .unwrap_err(); + assert!( + error.to_string().contains("exceeding the 7-byte limit"), + "unexpected error: {error}" + ); + assert_eq!(output, [7, 9]); } #[test] @@ -745,6 +983,22 @@ mod tests { .decompress(&compressed_data, &mut decompressed_data) .unwrap(); assert_eq!(input_data, decompressed_data.as_slice()); + + let mut exact = Vec::new(); + compressor + .decompress_exact(&compressed_data, &mut exact, input_data.len()) + .unwrap(); + assert_eq!(input_data, exact.as_slice()); + + let mut bounded = vec![3, 5]; + let error = compressor + .decompress_exact(&compressed_data, &mut bounded, input_data.len() - 1) + .unwrap_err(); + assert!( + error.to_string().contains("exceeds the 12-byte limit"), + "unexpected error: {error}" + ); + assert_eq!(bounded, [3, 5]); } } @@ -782,6 +1036,32 @@ mod tests { .decompress(&compressed_data, &mut decompressed_data) .unwrap(); assert_eq!(input_data, decompressed_data.as_slice()); + + let mut exact = Vec::new(); + compressor + .decompress_exact(&compressed_data, &mut exact, input_data.len()) + .unwrap(); + assert_eq!(input_data, exact.as_slice()); + } + + #[test] + fn test_lz4_rejects_untrusted_output_size() { + let compressor = Lz4BufferCompressor::default(); + let mut compressed_data = Vec::new(); + compressor + .compress(b"bounded", &mut compressed_data) + .unwrap(); + compressed_data[..4].copy_from_slice(&u32::MAX.to_le_bytes()); + + let mut output = vec![7, 9]; + let error = compressor + .decompress_exact(&compressed_data, &mut output, 7) + .unwrap_err(); + assert!( + error.to_string().contains("exceeding the 7-byte limit"), + "unexpected error: {error}" + ); + assert_eq!(output, [7, 9]); } #[test_log::test(tokio::test)] diff --git a/rust/lance-encoding/src/encodings/physical/constant.rs b/rust/lance-encoding/src/encodings/physical/constant.rs index c3fa16863f4..67f6aa6fa3d 100644 --- a/rust/lance-encoding/src/encodings/physical/constant.rs +++ b/rust/lance-encoding/src/encodings/physical/constant.rs @@ -3,13 +3,18 @@ //! Routines for compressing and decompressing constant-encoded data +#[cfg(test)] +use crate::compression::{BlockCompressor, block::validate_fixed_payload_len}; use crate::{ buffer::LanceBuffer, - compression::{BlockDecompressor, FixedPerValueDecompressor}, - data::{AllNullDataBlock, ConstantDataBlock, DataBlock, FixedWidthDataBlock}, + compression::{ + BlockDecompressor, BlockValueType, FixedPerValueDecompressor, require_no_block_payload, + }, + data::{AllNullDataBlock, BlockInfo, ConstantDataBlock, DataBlock, FixedWidthDataBlock}, + encodings::physical::try_vec_with_capacity, }; -use lance_core::Result; +use lance_core::{Error, Result}; /// A decompressor for constant-encoded data #[derive(Debug)] @@ -24,7 +29,7 @@ impl ConstantDecompressor { } impl BlockDecompressor for ConstantDecompressor { - fn decompress(&self, _data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, _data: Option, num_values: u64) -> Result { if let Some(scalar) = self.scalar.clone() { Ok(DataBlock::Constant(ConstantDataBlock { data: scalar, @@ -55,3 +60,134 @@ impl FixedPerValueDecompressor for ConstantDecompressor { .unwrap_or(0) } } + +/// Metadata-only fixed-width constant (or typed empty) block compressor. +#[cfg(test)] +#[derive(Debug)] +pub(crate) struct ConstantBlockCompressor { + value_type: BlockValueType, + value: Option, +} + +#[cfg(test)] +impl ConstantBlockCompressor { + pub(crate) fn new(value_type: BlockValueType, value: Option) -> Self { + Self { value_type, value } + } +} + +#[cfg(test)] +impl BlockCompressor for ConstantBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Constant block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Constant block compressor expects {}-bit values, got {}", + self.value_type.bits_per_value(), + data.bits_per_value + ))); + } + validate_fixed_payload_len( + &data.data, + self.value_type, + data.num_values, + "Constant block input", + )?; + match self.value { + None => { + if data.num_values != 0 { + return Err(Error::invalid_input( + "Typed empty block compressor received non-empty data", + )); + } + } + Some(expected) => { + if data.num_values == 0 { + return Err(Error::invalid_input( + "Constant block compressor received an empty sequence", + )); + } + macro_rules! check_values { + ($ty:ty) => { + for (index, actual) in + data.data.borrow_to_typed_slice::<$ty>().iter().enumerate() + { + if u64::from(*actual) != expected { + return Err(Error::invalid_input(format!( + "Constant block expects {expected}, got {actual} at index {index}" + ))); + } + } + }; + } + match self.value_type { + BlockValueType::UInt8 => check_values!(u8), + BlockValueType::UInt16 => check_values!(u16), + BlockValueType::UInt32 => check_values!(u32), + BlockValueType::UInt64 => check_values!(u64), + } + } + } + Ok(None) + } +} + +/// Metadata-only fixed-width constant (or typed empty) block decompressor. +#[derive(Debug)] +pub(crate) struct ConstantBlockDecompressor { + value_type: BlockValueType, + value: Option, +} + +impl ConstantBlockDecompressor { + pub(crate) fn new(value_type: BlockValueType, value: Option) -> Self { + Self { value_type, value } + } +} + +impl BlockDecompressor for ConstantBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + require_no_block_payload(data, "Constant block")?; + let output = match self.value { + None => { + if num_values != 0 { + return Err(Error::invalid_input(format!( + "Typed empty block cannot represent {num_values} values" + ))); + } + LanceBuffer::empty() + } + Some(value) => { + if num_values == 0 { + return Err(Error::invalid_input( + "Non-empty Constant descriptor cannot represent an empty block", + )); + } + macro_rules! repeat { + ($ty:ty) => {{ + let mut values = + try_vec_with_capacity::<$ty>(num_values, "Constant block output")?; + values.resize(num_values as usize, value as $ty); + LanceBuffer::reinterpret_vec(values) + }}; + } + match self.value_type { + BlockValueType::UInt8 => repeat!(u8), + BlockValueType::UInt16 => repeat!(u16), + BlockValueType::UInt32 => repeat!(u32), + BlockValueType::UInt64 => repeat!(u64), + } + } + }; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: output, + bits_per_value: self.value_type.bits_per_value(), + num_values, + block_info: BlockInfo::new(), + })) + } +} diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index 769f308ccec..e673d00dc39 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -3,10 +3,14 @@ use log::trace; +#[cfg(all(test, any(feature = "lz4", feature = "zstd")))] +use crate::compression::BlockCompressor; use crate::{ Result, buffer::LanceBuffer, - compression::MiniBlockDecompressor, + compression::{ + BlockDecompressor, BlockValueType, MiniBlockDecompressor, require_block_payload, + }, data::DataBlock, encodings::{ logical::primitive::miniblock::{ @@ -16,6 +20,88 @@ use crate::{ }, format::{ProtobufUtils21, pb21::CompressiveEncoding}, }; +use lance_core::Error; + +#[cfg(all(test, any(feature = "lz4", feature = "zstd")))] +pub(crate) fn compress_block( + compression: CompressionConfig, + payload: &[u8], +) -> Result { + let compressor = GeneralBufferCompressor::get_compressor(compression)?; + let mut compressed = Vec::new(); + compressor.compress(payload, &mut compressed)?; + Ok(LanceBuffer::from(compressed)) +} + +pub(crate) fn decompress_block_exact( + compression: CompressionConfig, + payload: &LanceBuffer, + expected_bytes: usize, +) -> Result { + let compressor = GeneralBufferCompressor::get_compressor(compression)?; + let mut decompressed = Vec::new(); + compressor.decompress_exact(payload, &mut decompressed, expected_bytes)?; + Ok(LanceBuffer::from(decompressed)) +} + +/// General-purpose block compressor that owns its child block compressor. +#[cfg(all(test, any(feature = "lz4", feature = "zstd")))] +#[derive(Debug)] +pub(crate) struct GeneralBlockCompressor { + child: Box, + compression: CompressionConfig, +} + +#[cfg(all(test, any(feature = "lz4", feature = "zstd")))] +impl GeneralBlockCompressor { + pub(crate) fn new(child: Box, compression: CompressionConfig) -> Self { + Self { child, compression } + } +} + +#[cfg(all(test, any(feature = "lz4", feature = "zstd")))] +impl BlockCompressor for GeneralBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let payload = self.child.compress(data)?.ok_or_else(|| { + Error::invalid_input("General block compression requires a payload-bearing child") + })?; + compress_block(self.compression, &payload).map(Some) + } +} + +/// Bounded general block decompressor used by generic fixed-width sequences. +#[derive(Debug)] +pub(crate) struct GenericGeneralBlockDecompressor { + child: Box, + compression: CompressionConfig, + value_type: BlockValueType, +} + +impl GenericGeneralBlockDecompressor { + pub(crate) fn new( + child: Box, + compression: CompressionConfig, + value_type: BlockValueType, + ) -> Self { + Self { + child, + compression, + value_type, + } + } +} + +impl BlockDecompressor for GenericGeneralBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "General block compression")?; + let expected_bytes = usize::try_from(num_values) + .ok() + .and_then(|values| values.checked_mul(self.value_type.bits_per_value() as usize / 8)) + .ok_or_else(|| Error::invalid_input("General block output length overflows usize"))?; + let decompressed = decompress_block_exact(self.compression, &data, expected_bytes)?; + self.child.decompress(Some(decompressed), num_values) + } +} /// A miniblock compressor that wraps another miniblock compressor and applies /// general-purpose compression (LZ4, Zstd) to the resulting buffers. @@ -29,22 +115,12 @@ impl GeneralMiniBlockCompressor { pub fn new(inner: Box, compression: CompressionConfig) -> Self { Self { inner, compression } } -} - -/// Minimum buffer size to consider for compression -const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024; - -use super::super::logical::primitive::miniblock::MiniBlockChunk; -impl MiniBlockCompressor for GeneralMiniBlockCompressor { - fn compress( + fn compress_inner( &self, - page: DataBlock, - context: MiniBlockCompressionContext, + inner_compressed: MiniBlockCompressed, + inner_encoding: CompressiveEncoding, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { - // First, compress with the inner compressor - let (inner_compressed, inner_encoding) = self.inner.compress(page, context)?; - // Return the original encoding without compression if there's no data or // the first buffer is not large enough if inner_compressed.data.is_empty() @@ -114,6 +190,22 @@ impl MiniBlockCompressor for GeneralMiniBlockCompressor { } } +/// Minimum buffer size to consider for compression +const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024; + +use super::super::logical::primitive::miniblock::MiniBlockChunk; + +impl MiniBlockCompressor for GeneralMiniBlockCompressor { + fn compress( + &self, + page: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + let (inner_compressed, inner_encoding) = self.inner.compress(page, context)?; + self.compress_inner(inner_compressed, inner_encoding) + } +} + /// A miniblock decompressor that first decompresses buffers using general-purpose /// compression (LZ4, Zstd) and then delegates to an inner miniblock decompressor. #[derive(Debug)] diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index 0127817d08e..4091e0ddefb 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -58,14 +58,23 @@ use arrow_buffer::{ArrowNativeType, ScalarBuffer}; use log::trace; use crate::buffer::LanceBuffer; -use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor}; +#[cfg(test)] +use crate::compression::block::{fixed_from_u64_values, visit_unsigned_values}; +use crate::compression::{ + BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, + block::{fixed_block, read_unsigned_values, validate_fixed_payload_len}, + require_block_payload, +}; use crate::data::DataBlock; use crate::data::{BlockInfo, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, }; -use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor}; +use crate::encodings::physical::{ + block::{BufferCompressor, CompressionConfig, GeneralBufferCompressor}, + try_vec_with_capacity, +}; use crate::format::ProtobufUtils21; use crate::format::pb21::CompressiveEncoding; @@ -872,7 +881,9 @@ impl RleEncoder { num_values: child_values, block_info: BlockInfo::default(), }); - let chunk_packed = BlockCompressor::compress(&compressor, block)?; + let chunk_packed = BlockCompressor::compress(&compressor, block)?.ok_or_else(|| { + Error::internal("RLE bitpacking child returned no payload".to_string()) + })?; let packed_size = u32::try_from(chunk_packed.len()).map_err(|_| { Error::invalid_input_source( format!( @@ -1107,7 +1118,7 @@ impl MiniBlockCompressor for RleEncoder { impl BlockCompressor for RleEncoder { // Block format: [8-byte header: values buffer size][values buffer][run_lengths buffer] - fn compress(&self, data: DataBlock) -> Result { + fn compress(&self, data: DataBlock) -> Result> { match data { DataBlock::FixedWidth(fixed_width) => { let num_values = fixed_width.num_values; @@ -1122,7 +1133,7 @@ impl BlockCompressor for RleEncoder { combined.extend_from_slice(&values_size.to_le_bytes()); combined.extend_from_slice(&all_buffers[0]); combined.extend_from_slice(&all_buffers[1]); - Ok(LanceBuffer::from(combined)) + Ok(Some(LanceBuffer::from(combined))) } _ => Err(Error::invalid_input_source( "RLE encoding only supports FixedWidth data blocks".into(), @@ -1149,6 +1160,9 @@ pub(crate) struct RleChildDecompressor { #[derive(Debug)] enum RleChildDecompressorInner { Flat, + General { + compressor: Box, + }, Block { decompressor: Box, requires_num_values: bool, @@ -1177,13 +1191,22 @@ impl RleChildDecompressor { } } + pub(crate) fn general(bits_per_value: u64, compression: CompressionConfig) -> Result { + Ok(Self { + bits_per_value, + inner: RleChildDecompressorInner::General { + compressor: GeneralBufferCompressor::get_compressor(compression)?, + }, + }) + } + pub(crate) fn bits_per_value(&self) -> u64 { self.bits_per_value } pub(crate) fn requires_num_values(&self) -> bool { match &self.inner { - RleChildDecompressorInner::Flat => false, + RleChildDecompressorInner::Flat | RleChildDecompressorInner::General { .. } => false, RleChildDecompressorInner::Block { requires_num_values, .. @@ -1199,10 +1222,66 @@ impl RleChildDecompressor { &self, data: LanceBuffer, num_values: Option, + max_num_values: u64, label: &str, ) -> Result { match &self.inner { RleChildDecompressorInner::Flat => Ok(data), + RleChildDecompressorInner::General { compressor } => { + let bytes_per_value = usize::try_from(self.bits_per_value / 8).map_err(|_| { + Error::invalid_input_source( + format!( + "RLE {label} child bit width is too large: {}", + self.bits_per_value + ) + .into(), + ) + })?; + let max_output_bytes = usize::try_from(max_num_values) + .ok() + .and_then(|num_values| num_values.checked_mul(bytes_per_value)) + .ok_or_else(|| { + Error::invalid_input_source( + format!("RLE {label} child maximum payload length overflows usize") + .into(), + ) + })?; + let mut decompressed = Vec::new(); + if let Some(num_values) = num_values { + if num_values > max_num_values { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child expects {num_values} values, exceeding the maximum run count {max_num_values}" + ) + .into(), + )); + } + let expected_bytes = usize::try_from(num_values) + .ok() + .and_then(|num_values| num_values.checked_mul(bytes_per_value)) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "RLE {label} child expected payload length overflows usize" + ) + .into(), + ) + })?; + compressor.decompress_exact(&data, &mut decompressed, expected_bytes)?; + } else { + compressor.decompress_bounded(&data, &mut decompressed, max_output_bytes)?; + } + if bytes_per_value == 0 || !decompressed.len().is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decompressed to {} bytes, not divisible by {bytes_per_value}", + decompressed.len() + ) + .into(), + )); + } + Ok(LanceBuffer::from(decompressed)) + } RleChildDecompressorInner::Block { decompressor, requires_num_values, @@ -1216,7 +1295,7 @@ impl RleChildDecompressor { } else { num_values.unwrap_or(0) }; - let decoded = decompressor.decompress(data, num_values)?; + let decoded = decompressor.decompress(Some(data), num_values)?; self.extract_fixed_width(decoded, num_values, label) } } @@ -1322,7 +1401,7 @@ impl RleDecompressor { let values_buffer = data_iter.next().unwrap(); let lengths_buffer = data_iter.next().unwrap(); let (values_buffer, lengths_buffer) = - self.decode_child_buffers(values_buffer, lengths_buffer)?; + self.decode_child_buffers(values_buffer, lengths_buffer, num_values)?; let decoded_data = match self.bits_per_value { 8 => self.decode_generic::( @@ -1372,6 +1451,7 @@ impl RleDecompressor { &self, values_buffer: LanceBuffer, lengths_buffer: LanceBuffer, + max_num_runs: u64, ) -> Result<(LanceBuffer, LanceBuffer)> { let values_requires_num_runs = self.values.requires_num_values(); let lengths_requires_num_runs = self.run_lengths.requires_num_values(); @@ -1382,31 +1462,38 @@ impl RleDecompressor { } if values_requires_num_runs { - let lengths_buffer = self - .run_lengths - .decode(lengths_buffer, None, "run lengths")?; + let lengths_buffer = + self.run_lengths + .decode(lengths_buffer, None, max_num_runs, "run lengths")?; let num_runs = Self::num_child_values( &lengths_buffer, self.run_lengths.bits_per_value(), "run lengths", )?; - let values_buffer = self - .values - .decode(values_buffer, Some(num_runs), "values")?; + let values_buffer = + self.values + .decode(values_buffer, Some(num_runs), max_num_runs, "values")?; Ok((values_buffer, lengths_buffer)) } else if lengths_requires_num_runs { - let values_buffer = self.values.decode(values_buffer, None, "values")?; + let values_buffer = self + .values + .decode(values_buffer, None, max_num_runs, "values")?; let num_runs = Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?; - let lengths_buffer = - self.run_lengths - .decode(lengths_buffer, Some(num_runs), "run lengths")?; + let lengths_buffer = self.run_lengths.decode( + lengths_buffer, + Some(num_runs), + max_num_runs, + "run lengths", + )?; Ok((values_buffer, lengths_buffer)) } else { - let values_buffer = self.values.decode(values_buffer, None, "values")?; - let lengths_buffer = self - .run_lengths - .decode(lengths_buffer, None, "run lengths")?; + let values_buffer = self + .values + .decode(values_buffer, None, max_num_runs, "values")?; + let lengths_buffer = + self.run_lengths + .decode(lengths_buffer, None, max_num_runs, "run lengths")?; Ok((values_buffer, lengths_buffer)) } } @@ -1559,7 +1646,8 @@ impl MiniBlockDecompressor for RleDecompressor { } impl BlockDecompressor for RleDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "RLE")?; let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?; self.decode_data(vec![values_buffer, lengths_buffer], num_values, false) } @@ -1829,7 +1917,7 @@ impl RleDecompressor { } let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?; let (values_buffer, lengths_buffer) = - self.decode_child_buffers(values_buffer, lengths_buffer)?; + self.decode_child_buffers(values_buffer, lengths_buffer, num_values)?; RleRuns::try_new( values_buffer, lengths_buffer, @@ -1839,16 +1927,422 @@ impl RleDecompressor { } } +// Generic block codec support. + +pub(crate) const BLOCK_FRAME_BYTES: u64 = 8; + +/// Generic RLE compressor that owns both child compressors. +#[cfg(test)] +#[derive(Debug)] +pub(crate) struct BlockRleCompressor { + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, +} + +#[cfg(test)] +impl BlockRleCompressor { + pub(crate) fn new( + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, + ) -> Self { + Self { + value_type, + run_length_type, + values, + run_lengths, + } + } +} + +#[cfg(test)] +impl BlockCompressor for BlockRleCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "RLE block compression requires fixed-width data", + )); + }; + let (run_values, run_lengths) = + materialize_block(&data, self.value_type, self.run_length_type)?; + let values_payload = self.values.compress(DataBlock::FixedWidth(run_values))?; + let lengths_payload = self + .run_lengths + .compress(DataBlock::FixedWidth(run_lengths))?; + + if values_payload.is_none() && lengths_payload.is_none() { + return Ok(None); + } + + let values_payload = values_payload.unwrap_or_else(LanceBuffer::empty); + let lengths_payload = lengths_payload.unwrap_or_else(LanceBuffer::empty); + let mut output = try_block_frame(values_payload.len(), lengths_payload.len())?; + output.extend_from_slice(&(values_payload.len() as u64).to_le_bytes()); + output.extend_from_slice(&values_payload); + output.extend_from_slice(&lengths_payload); + Ok(Some(LanceBuffer::from(output))) + } +} + +#[cfg(test)] +fn try_block_frame(values_payload_bytes: usize, lengths_payload_bytes: usize) -> Result> { + let capacity = (BLOCK_FRAME_BYTES as usize) + .checked_add(values_payload_bytes) + .and_then(|capacity| capacity.checked_add(lengths_payload_bytes)) + .ok_or_else(|| Error::invalid_input("RLE frame length overflows usize"))?; + let mut output = Vec::new(); + output.try_reserve_exact(capacity).map_err(|error| { + Error::invalid_input(format!( + "RLE could not reserve {capacity} frame bytes: {error}" + )) + })?; + Ok(output) +} + +/// Metadata-only run lengths used to recover the RLE child cardinality. +#[derive(Debug, Clone, Copy)] +pub(crate) enum MetadataRunLengths { + Constant(u64), +} + +impl MetadataRunLengths { + fn infer_run_count(self, num_values: u64) -> Result { + match self { + Self::Constant(run_length) => { + if run_length == 0 { + return Err(Error::invalid_input("RLE run lengths must be positive")); + } + if !num_values.is_multiple_of(run_length) { + return Err(Error::invalid_input(format!( + "RLE constant run length {run_length} does not divide {num_values} values" + ))); + } + let run_count = num_values / run_length; + if run_count == 0 { + return Err(Error::invalid_input( + "RLE metadata describes zero runs for a non-empty sequence", + )); + } + Ok(run_count) + } + } + } +} + +/// The descriptor-derived source of the RLE run count. +#[derive(Debug, Clone, Copy)] +pub(crate) enum BlockRunCount { + Metadata(MetadataRunLengths), + ValuesPayload, + RunLengthsPayload, +} + +/// Generic RLE decompressor that owns both child decompressors. +#[derive(Debug)] +pub(crate) struct BlockRleDecompressor { + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, + values_have_payload: bool, + run_lengths_have_payload: bool, + run_count: BlockRunCount, +} + +impl BlockRleDecompressor { + pub(crate) fn new( + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, + values_have_payload: bool, + run_lengths_have_payload: bool, + run_count: BlockRunCount, + ) -> Self { + Self { + value_type, + run_length_type, + values, + run_lengths, + values_have_payload, + run_lengths_have_payload, + run_count, + } + } +} + +impl BlockDecompressor for BlockRleDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + if num_values == 0 { + return Err(Error::invalid_input( + "RLE cannot represent an empty block sequence", + )); + } + let (values_payload, lengths_payload) = if self.values_have_payload + || self.run_lengths_have_payload + { + let data = require_block_payload(data, "RLE")?; + if data.len() < BLOCK_FRAME_BYTES as usize { + return Err(Error::invalid_input(format!( + "RLE payload has {} bytes, shorter than its {BLOCK_FRAME_BYTES}-byte header", + data.len() + ))); + } + let values_size = u64::from_le_bytes( + data[..BLOCK_FRAME_BYTES as usize] + .try_into() + .expect("RLE header length was checked"), + ); + let values_size = usize::try_from(values_size).map_err(|_| { + Error::invalid_input("RLE values payload length does not fit usize") + })?; + let values_start = BLOCK_FRAME_BYTES as usize; + let lengths_start = values_start + .checked_add(values_size) + .ok_or_else(|| Error::invalid_input("RLE values payload end overflows usize"))?; + if lengths_start > data.len() { + return Err(Error::invalid_input(format!( + "RLE values payload ends at {lengths_start}, beyond {} bytes", + data.len() + ))); + } + let values_payload = data.slice_with_length(values_start, values_size); + let lengths_payload = data.slice_with_length(lengths_start, data.len() - lengths_start); + if !self.values_have_payload && !values_payload.is_empty() { + return Err(Error::invalid_input(format!( + "Metadata-only RLE values child has {} framed payload bytes", + values_payload.len() + ))); + } + if !self.run_lengths_have_payload && !lengths_payload.is_empty() { + return Err(Error::invalid_input(format!( + "Metadata-only RLE run-length child has {} framed payload bytes", + lengths_payload.len() + ))); + } + ( + self.values_have_payload.then_some(values_payload), + self.run_lengths_have_payload.then_some(lengths_payload), + ) + } else { + if data.is_some() { + return Err(Error::invalid_input("Metadata-only RLE expects no payload")); + } + (None, None) + }; + + let run_count = match self.run_count { + BlockRunCount::Metadata(metadata) => metadata.infer_run_count(num_values)?, + BlockRunCount::ValuesPayload => infer_flat_run_count( + values_payload.as_ref().ok_or_else(|| { + Error::invalid_input("RLE values payload is required to infer the run count") + })?, + self.value_type, + num_values, + "RLE values", + )?, + BlockRunCount::RunLengthsPayload => infer_flat_run_count( + lengths_payload.as_ref().ok_or_else(|| { + Error::invalid_input( + "RLE run-length payload is required to infer the run count", + ) + })?, + self.run_length_type, + num_values, + "RLE run lengths", + )?, + }; + if run_count == 0 { + return Err(Error::invalid_input( + "RLE payload contains zero runs for a non-empty sequence", + )); + } + let values = self.values.decompress(values_payload, run_count)?; + let run_lengths = self.run_lengths.decompress(lengths_payload, run_count)?; + expand_block( + values, + run_lengths, + self.value_type, + self.run_length_type, + num_values, + run_count, + ) + } +} + +fn infer_flat_run_count( + payload: &LanceBuffer, + value_type: BlockValueType, + max_run_count: u64, + label: &str, +) -> Result { + let bytes_per_value = value_type.bytes_per_value(); + if !payload.len().is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input(format!( + "{label} payload has {} bytes, not divisible by {bytes_per_value}", + payload.len() + ))); + } + let run_count = (payload.len() / bytes_per_value) as u64; + if run_count > max_run_count { + return Err(Error::invalid_input(format!( + "{label} payload contains {run_count} runs, exceeding the {max_run_count}-run limit" + ))); + } + Ok(run_count) +} + +pub(crate) fn expand_block( + values: DataBlock, + run_lengths: DataBlock, + value_type: BlockValueType, + run_length_type: BlockValueType, + num_values: u64, + run_count: u64, +) -> Result { + let DataBlock::FixedWidth(values) = values else { + return Err(Error::invalid_input( + "RLE values decoded to a non fixed-width block", + )); + }; + let DataBlock::FixedWidth(run_lengths) = run_lengths else { + return Err(Error::invalid_input( + "RLE run lengths decoded to a non fixed-width block", + )); + }; + if values.num_values != run_count + || values.bits_per_value != value_type.bits_per_value() + || run_lengths.num_values != run_count + || run_lengths.bits_per_value != run_length_type.bits_per_value() + { + return Err(Error::invalid_input( + "RLE child cardinality or bit width does not match its descriptor", + )); + } + validate_fixed_payload_len(&values.data, value_type, run_count, "RLE values")?; + validate_fixed_payload_len( + &run_lengths.data, + run_length_type, + run_count, + "RLE run lengths", + )?; + + let lengths = read_unsigned_values(&run_lengths, run_length_type)?; + let mut total = 0_u64; + for (index, length) in lengths.iter().enumerate() { + if *length == 0 { + return Err(Error::invalid_input(format!( + "RLE run length at index {index} is zero" + ))); + } + total = total.checked_add(*length).ok_or_else(|| { + Error::invalid_input(format!("RLE run length sum overflows at index {index}")) + })?; + } + if total != num_values { + return Err(Error::invalid_input(format!( + "RLE run lengths sum to {total}, expected {num_values}" + ))); + } + let lengths = lengths + .into_iter() + .map(|length| { + usize::try_from(length) + .map_err(|_| Error::invalid_input("RLE run length does not fit usize")) + }) + .collect::>>()?; + + let run_values = read_unsigned_values(&values, value_type)?; + let output = match value_type { + BlockValueType::UInt8 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value as u8, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt16 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value as u16, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt32 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value as u32, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt64 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + }; + Ok(fixed_block(value_type, num_values, output)) +} + +#[cfg(test)] +pub(crate) fn materialize_block( + data: &FixedWidthDataBlock, + value_type: BlockValueType, + run_length_type: BlockValueType, +) -> Result<(FixedWidthDataBlock, FixedWidthDataBlock)> { + if data.num_values == 0 { + return Err(Error::invalid_input( + "RLE cannot materialize an empty sequence", + )); + } + let max_run_length = run_length_type.max_value(); + let mut run_values = Vec::new(); + let mut run_lengths = Vec::new(); + let mut current = None; + let mut length = 0_u64; + visit_unsigned_values(data, value_type, |value| { + match current { + Some(current_value) if value == current_value && length < max_run_length => { + length += 1; + } + Some(current_value) => { + run_values.push(current_value); + run_lengths.push(length); + current = Some(value); + length = 1; + } + None => { + current = Some(value); + length = 1; + } + } + Ok(()) + })?; + run_values.push(current.expect("non-empty input has a run value")); + run_lengths.push(length); + Ok(( + fixed_from_u64_values(&run_values, value_type, "RLE values")?, + fixed_from_u64_values(&run_lengths, run_length_type, "RLE run lengths")?, + )) +} + #[cfg(test)] mod tests { use std::sync::Arc; use super::*; - use crate::compression::{ - DecompressionStrategy, DefaultDecompressionStrategy, create_rle_decompressor, - }; + #[cfg(any(feature = "lz4", feature = "zstd"))] + use crate::compression::create_rle_decompressor; + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] + use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; use crate::data::DataBlock; use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES; + #[cfg(any(feature = "lz4", feature = "zstd"))] use crate::encodings::physical::block::{CompressionConfig, CompressionScheme}; use crate::{ buffer::LanceBuffer, @@ -1857,13 +2351,6 @@ mod tests { use arrow_array::Int32Array; use rstest::rstest; - fn compress_miniblock( - compressor: &dyn MiniBlockCompressor, - data: DataBlock, - ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { - compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) - } - fn expand_u16_runs(runs: &RleRuns) -> Vec { let mut expanded = Vec::with_capacity(runs.num_values()); for (value, length) in runs.iter() { @@ -1886,11 +2373,16 @@ mod tests { num_values, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); - let eager = - BlockDecompressor::decompress(&RleDecompressor::new(16), frame.clone(), num_values) - .unwrap(); + let eager = BlockDecompressor::decompress( + &RleDecompressor::new(16), + Some(frame.clone()), + num_values, + ) + .unwrap(); let DataBlock::FixedWidth(eager) = eager else { panic!("expected fixed-width block"); }; @@ -1963,7 +2455,9 @@ mod tests { num_values, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); let (values, lengths) = parse_rle_block_frame(&frame).unwrap(); let compression = test_general_compression(); @@ -2005,7 +2499,9 @@ mod tests { num_values, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); let runs = RleDecompressor::new(16) .decode_u16_runs(frame, num_values) .unwrap(); @@ -2026,7 +2522,9 @@ mod tests { num_values: n, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); let runs = RleDecompressor::new(16).decode_u16_runs(frame, n).unwrap(); assert_eq!( runs.coalesced_runs() as u64, @@ -2036,6 +2534,19 @@ mod tests { assert_eq!(expand_u16_runs(&runs), alternating); } + fn compress_miniblock( + compressor: &dyn MiniBlockCompressor, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) + } + + #[test] + fn block_frame_capacity_overflow_is_fallible() { + let error = try_block_frame(usize::MAX, 1).unwrap_err(); + assert!(error.to_string().contains("frame length overflows usize")); + } + // ========== Core Functionality Tests ========== #[test] @@ -2205,7 +2716,9 @@ mod tests { block_info: BlockInfo::default(), }); let bitpacked_run_lengths = - BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block).unwrap(); + BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block) + .unwrap() + .unwrap(); let encoding = ProtobufUtils21::rle( ProtobufUtils21::flat(32, None), ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)), @@ -2249,6 +2762,7 @@ mod tests { } } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn repeating_runs(num_runs: usize, run_length: usize) -> Vec { let mut values = Vec::with_capacity(num_runs * run_length); for run in 0..num_runs { @@ -2257,6 +2771,7 @@ mod tests { values } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { match encoding.compression.as_ref().unwrap() { crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle, @@ -2264,6 +2779,7 @@ mod tests { } } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) { match decoded { DataBlock::FixedWidth(block) => { @@ -2371,6 +2887,7 @@ mod tests { assert_eq!(decoded, expected); } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn decompress_i32_chunks( compressed: &MiniBlockCompressed, encoding: &CompressiveEncoding, @@ -2693,8 +3210,9 @@ mod tests { payload.extend_from_slice(&values); payload.extend_from_slice(&lengths); - let error = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(payload), 5) - .unwrap_err(); + let error = + BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(payload)), 5) + .unwrap_err(); assert!(matches!(&error, Error::InvalidInput { .. })); assert!( error @@ -3140,13 +3658,97 @@ mod tests { } // ========== Block Related tests ========== + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn assert_rle_general_children_reject_oversized_output( + compression: CompressionConfig, + payload: Vec, + ) { + let max_num_values = 2; + + let values_general = RleChildDecompressor::general(32, compression).unwrap(); + let values_decompressor = RleDecompressor::with_child_decompressors( + 32, + RunLengthWidth::U8, + values_general, + RleChildDecompressor::flat(8), + ); + let error = MiniBlockDecompressor::decompress( + &values_decompressor, + vec![ + LanceBuffer::from(payload.clone()), + LanceBuffer::from(vec![2_u8]), + ], + max_num_values, + ) + .unwrap_err(); + assert!( + error.to_string().contains("exceed"), + "unexpected values-child error: {error}" + ); + + let lengths_general = RleChildDecompressor::general(8, compression).unwrap(); + let lengths_decompressor = RleDecompressor::with_child_decompressors( + 32, + RunLengthWidth::U8, + RleChildDecompressor::flat(32), + lengths_general, + ); + let error = MiniBlockDecompressor::decompress( + &lengths_decompressor, + vec![ + LanceBuffer::reinterpret_vec(vec![7_u32]), + LanceBuffer::from(payload), + ], + max_num_values, + ) + .unwrap_err(); + assert!( + error.to_string().contains("exceed"), + "unexpected run-lengths-child error: {error}" + ); + } + + #[cfg(feature = "lz4")] + #[test] + fn test_rle_general_children_bound_lz4_output() { + assert_rle_general_children_reject_oversized_output( + CompressionConfig::new(CompressionScheme::Lz4, None), + u32::MAX.to_le_bytes().to_vec(), + ); + } + + #[cfg(feature = "zstd")] + #[test] + fn test_rle_general_children_bound_length_prefixed_zstd_output() { + assert_rle_general_children_reject_oversized_output( + CompressionConfig::new(CompressionScheme::Zstd, Some(0)), + u64::MAX.to_le_bytes().to_vec(), + ); + } + + #[cfg(feature = "zstd")] + #[test] + fn test_rle_general_children_bound_raw_zstd_output() { + use std::io::Write; + + let mut payload = Vec::new(); + let mut encoder = ::zstd::Encoder::new(&mut payload, 0).unwrap(); + encoder.write_all(&[0; 1024]).unwrap(); + encoder.finish().unwrap(); + + assert_rle_general_children_reject_oversized_output( + CompressionConfig::new(CompressionScheme::Zstd, Some(0)), + payload, + ); + } + #[test] fn test_block_decompressor_rejects_overflowing_values_size() { let decompressor = RleDecompressor::new(32); let mut data = Vec::new(); data.extend_from_slice(&u64::MAX.to_le_bytes()); - let result = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(data), 1); + let result = BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(data)), 1); assert!(result.is_err()); assert!( result @@ -3159,8 +3761,11 @@ mod tests { #[test] fn test_block_decompressor_too_small() { let decompressor = RleDecompressor::new(32); - let result = - BlockDecompressor::decompress(&decompressor, LanceBuffer::from(vec![1, 2, 3]), 10); + let result = BlockDecompressor::decompress( + &decompressor, + Some(LanceBuffer::from(vec![1, 2, 3])), + 10, + ); assert!(result.is_err()); assert!( result @@ -3176,7 +3781,9 @@ mod tests { let data = vec![1i32, 1, 1]; let array = Int32Array::from(data); - let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)) + .unwrap() + .unwrap(); // Verify header format: first 8 bytes should be values_size as u64 assert!(compressed.len() >= 8); diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index 8b7c385f601..3b7f2365393 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -5,7 +5,8 @@ use arrow_buffer::{BooleanBufferBuilder, bit_util}; use crate::buffer::LanceBuffer; use crate::compression::{ - BlockCompressor, BlockDecompressor, FixedPerValueDecompressor, MiniBlockDecompressor, + BlockCompressor, BlockDecompressor, BlockValueType, FixedPerValueDecompressor, + MiniBlockDecompressor, require_block_payload, }; use crate::data::{ BlockInfo, DataBlock, FixedSizeListBlock, FixedWidthDataBlock, NullableDataBlock, @@ -458,15 +459,60 @@ impl ValueEncoder { } impl BlockCompressor for ValueEncoder { - fn compress(&self, data: DataBlock) -> Result { + fn compress(&self, data: DataBlock) -> Result> { let data = match data { DataBlock::FixedWidth(fixed_width) => fixed_width.data, - _ => unimplemented!( - "Cannot compress block of type {} with ValueEncoder", - data.name() - ), + _ => { + return Err(Error::invalid_input(format!( + "ValueEncoder cannot compress a {} block", + data.name() + ))); + } }; - Ok(data) + Ok(Some(data)) + } +} + +/// Flat fixed-width block compressor used by the generic block selector. +#[cfg(test)] +#[derive(Debug)] +pub(crate) struct FixedWidthBlockCompressor { + value_type: BlockValueType, +} + +#[cfg(test)] +impl FixedWidthBlockCompressor { + pub(crate) fn new(value_type: BlockValueType) -> Self { + Self { value_type } + } +} + +#[cfg(test)] +impl BlockCompressor for FixedWidthBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Flat block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Flat block compressor expects {}-bit values, got {}", + self.value_type.bits_per_value(), + data.bits_per_value + ))); + } + let expected = usize::try_from(data.num_values) + .ok() + .and_then(|values| values.checked_mul(self.value_type.bits_per_value() as usize / 8)) + .ok_or_else(|| Error::invalid_input("Flat block payload length overflows usize"))?; + if data.data.len() != expected { + return Err(Error::invalid_input(format!( + "Flat block input has {} bytes, expected {expected}", + data.data.len() + ))); + } + Ok(Some(data.data)) } } @@ -575,13 +621,48 @@ impl ValueDecompressor { } impl BlockDecompressor for ValueDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Flat block")?; let block = self.buffer_to_block(data, num_values); assert_eq!(block.num_values(), num_values); Ok(block) } } +/// Flat fixed-width block decompressor with fallible payload validation. +#[derive(Debug)] +pub(crate) struct FixedWidthBlockDecompressor { + value_type: BlockValueType, +} + +impl FixedWidthBlockDecompressor { + pub(crate) fn new(value_type: BlockValueType) -> Self { + Self { value_type } + } +} + +impl BlockDecompressor for FixedWidthBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Flat block")?; + let expected = usize::try_from(num_values) + .ok() + .and_then(|values| values.checked_mul(self.value_type.bits_per_value() as usize / 8)) + .ok_or_else(|| Error::invalid_input("Flat block payload length overflows usize"))?; + if data.len() != expected { + return Err(Error::invalid_input(format!( + "Flat block payload has {} bytes, expected {expected}", + data.len() + ))); + } + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.value_type.bits_per_value(), + num_values, + data, + block_info: BlockInfo::new(), + })) + } +} + impl MiniBlockDecompressor for ValueDecompressor { fn decompress(&self, data: Vec, num_values: u64) -> Result { let num_items = num_values * self.items_per_value; @@ -779,7 +860,9 @@ mod tests { encodings::{ logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::{MiniBlockCompressionContext, MiniBlockCompressor}, + miniblock::{ + MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, + }, }, physical::value::ValueDecompressor, }, @@ -793,8 +876,14 @@ mod tests { use super::ValueEncoder; - fn miniblock_context() -> MiniBlockCompressionContext { - MiniBlockCompressionContext::new(0, true, true) + fn compress_miniblock( + compressor: &dyn MiniBlockCompressor, + data: DataBlock, + ) -> lance_core::Result<( + MiniBlockCompressed, + crate::format::pb21::CompressiveEncoding, + )> { + compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) } const PRIMITIVE_TYPES: &[DataType] = &[ @@ -977,8 +1066,7 @@ mod tests { let starting_data = DataBlock::from_array(sample_list.clone()); let encoder = ValueEncoder::default(); - let (data, compression) = - MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); + let (data, compression) = compress_miniblock(&encoder, starting_data).unwrap(); assert_eq!(data.num_values, 3); assert_eq!(data.data.len(), 3); @@ -1039,7 +1127,7 @@ mod tests { let starting_data = DataBlock::from_array(array); let encoder = ValueEncoder::default(); - let result = MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()); + let result = compress_miniblock(&encoder, starting_data); let err = result.expect_err("wide values should not be encodable as miniblock"); assert!( @@ -1151,8 +1239,7 @@ mod tests { ); let encoder = ValueEncoder::default(); - let (data, compression) = - MiniBlockCompressor::compress(&encoder, starting_data, miniblock_context()).unwrap(); + let (data, compression) = compress_miniblock(&encoder, starting_data).unwrap(); let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { panic!() diff --git a/rust/lance-encoding/tests/compression_strategy.rs b/rust/lance-encoding/tests/compression_strategy.rs new file mode 100644 index 00000000000..23e3a70dddf --- /dev/null +++ b/rust/lance-encoding/tests/compression_strategy.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_schema::{DataType, Field as ArrowField}; +use lance_core::{Result, datatypes::Field}; +use lance_encoding::{ + buffer::LanceBuffer, + compression::{BlockCompressor, CompressionStrategy, DefaultCompressionStrategy}, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + format::ProtobufUtils21, +}; + +#[derive(Debug)] +struct IdentityBlockCompressor; + +impl BlockCompressor for IdentityBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + panic!("test compressor only accepts fixed-width data"); + }; + Ok(Some(data.data)) + } +} + +#[derive(Debug, Default)] +struct CustomCompressionStrategy { + fallback: DefaultCompressionStrategy, +} + +impl CompressionStrategy for CustomCompressionStrategy { + fn create_block_compressor( + &self, + _field: &Field, + data: &DataBlock, + ) -> Result<( + Box, + lance_encoding::format::pb21::CompressiveEncoding, + )> { + let DataBlock::FixedWidth(data) = data else { + panic!("test strategy only accepts fixed-width data"); + }; + Ok(( + Box::new(IdentityBlockCompressor), + ProtobufUtils21::flat(data.bits_per_value, None), + )) + } + + fn create_per_value( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + self.fallback.create_per_value(field, data) + } + + fn create_miniblock_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + self.fallback.create_miniblock_compressor(field, data) + } +} + +#[test] +fn public_strategy_returns_the_frozen_block_compressor() { + let field = Field::try_from(&ArrowField::new("values", DataType::UInt32, false)).unwrap(); + let values = vec![3_u32, 5, 8, 13]; + let data = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(values.clone()), + bits_per_value: 32, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + }); + let strategy: Box = Box::new(CustomCompressionStrategy::default()); + + let (compressor, _) = strategy.create_block_compressor(&field, &data).unwrap(); + let payload = compressor.compress(data).unwrap().unwrap(); + assert_eq!(payload.borrow_to_typed_slice::().as_ref(), values); +} From 166e42e496e9c84602a680d03059a4088b501b5a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 31 Jul 2026 00:17:09 +0800 Subject: [PATCH 3/5] refactor(encoding): adapt block call sites after merge --- rust/lance-encoding/benches/decoder.rs | 2 +- rust/lance-encoding/src/compression.rs | 7 ------- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/rust/lance-encoding/benches/decoder.rs b/rust/lance-encoding/benches/decoder.rs index abfad51b1b4..a2b3d64e0f2 100644 --- a/rust/lance-encoding/benches/decoder.rs +++ b/rust/lance-encoding/benches/decoder.rs @@ -639,7 +639,7 @@ where #[cfg(feature = "bitpacking")] fn typed_view_unchunk(buffer: LanceBuffer, uncompressed_bits: u64, num_values: u64) -> DataBlock { InlineBitpacking::new(uncompressed_bits) - .decompress(buffer, num_values) + .decompress(Some(buffer), num_values) .unwrap() } diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 301e6048fc2..dd44ecf0a52 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -1476,13 +1476,6 @@ mod tests { strategy.create_block_compressor(field, data).unwrap() } - fn miniblock_context() - -> crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext { - crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext::new( - 0, true, true, - ) - } - fn create_fixed_width_block_with_stats( bits_per_value: u64, num_values: u64, From f50f95e30cb60bd01be5589c28cf1a0dbb4f46fc Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 31 Jul 2026 14:41:23 +0800 Subject: [PATCH 4/5] refactor(encoding): adapt block contract to version strategies --- rust/lance-encoding/src/compression.rs | 5 ++++- .../logical/primitive/sparse/writer.rs | 13 ++++++----- .../tests/compression_strategy.rs | 22 +++++++++---------- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 94e78a24e89..f60962883a9 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -212,6 +212,9 @@ fn rle_beats_raw_and_bitpacking( encoded_bytes: u128, raw_bytes: u128, ) -> bool { + #[cfg(not(feature = "bitpacking"))] + let _ = data; + if encoded_bytes >= raw_bytes { return false; } @@ -1467,7 +1470,7 @@ mod tests { } fn selected_block_codec( - strategy: &dyn CompressionStrategy, + strategy: &Arc, field: &Field, data: &DataBlock, ) -> (Box, CompressiveEncoding) { diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index 312e616b2f4..a5811a74b1f 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -1575,9 +1575,13 @@ mod tests { )), ); let metadata = sparse_metadata(); - let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, metadata.clone()) - .await - .unwrap(); + let pages = encode_pages( + array.clone(), + TestEncoding::StructuralSparse, + metadata.clone(), + ) + .await + .unwrap(); assert_eq!(pages.len(), 1); let sparse = sparse_layout(&pages[0]); assert_eq!(sparse.num_buffers, 1); @@ -1604,8 +1608,7 @@ mod tests { )); let cases = TestCases::default() - .with_min_file_version(LanceFileVersion::V2_3) - .with_max_file_version(LanceFileVersion::V2_3) + .with_encoding(TestEncoding::StructuralSparse) .with_page_sizes(vec![1]) .with_range(1..17) .with_indices(vec![0, 7, (num_values * 2 - 1) as u64]); diff --git a/rust/lance-encoding/tests/compression_strategy.rs b/rust/lance-encoding/tests/compression_strategy.rs index 23e3a70dddf..707052cf585 100644 --- a/rust/lance-encoding/tests/compression_strategy.rs +++ b/rust/lance-encoding/tests/compression_strategy.rs @@ -5,7 +5,7 @@ use arrow_schema::{DataType, Field as ArrowField}; use lance_core::{Result, datatypes::Field}; use lance_encoding::{ buffer::LanceBuffer, - compression::{BlockCompressor, CompressionStrategy, DefaultCompressionStrategy}, + compression::{BlockCompressor, CompressionStrategy}, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, format::ProtobufUtils21, @@ -23,10 +23,8 @@ impl BlockCompressor for IdentityBlockCompressor { } } -#[derive(Debug, Default)] -struct CustomCompressionStrategy { - fallback: DefaultCompressionStrategy, -} +#[derive(Debug)] +struct CustomCompressionStrategy; impl CompressionStrategy for CustomCompressionStrategy { fn create_block_compressor( @@ -48,18 +46,18 @@ impl CompressionStrategy for CustomCompressionStrategy { fn create_per_value( &self, - field: &Field, - data: &DataBlock, + _field: &Field, + _data: &DataBlock, ) -> Result> { - self.fallback.create_per_value(field, data) + unreachable!("test only exercises block compression") } fn create_miniblock_compressor( &self, - field: &Field, - data: &DataBlock, + _field: &Field, + _data: &DataBlock, ) -> Result> { - self.fallback.create_miniblock_compressor(field, data) + unreachable!("test only exercises block compression") } } @@ -73,7 +71,7 @@ fn public_strategy_returns_the_frozen_block_compressor() { num_values: values.len() as u64, block_info: BlockInfo::default(), }); - let strategy: Box = Box::new(CustomCompressionStrategy::default()); + let strategy: Box = Box::new(CustomCompressionStrategy); let (compressor, _) = strategy.create_block_compressor(&field, &data).unwrap(); let payload = compressor.compress(data).unwrap().unwrap(); From de91c0a079c0e4cab38c0547df88621cf3841471 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 31 Jul 2026 14:41:35 +0800 Subject: [PATCH 5/5] fix(encoding): preserve zero-width bitpacking blocks --- .../src/encodings/physical/bitpacking.rs | 26 ++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index 8200b8054a4..6bebe80077a 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -340,7 +340,7 @@ impl BlockCompressor for InlineBitpacking { max_value = max_value.max(value); Ok(()) })?; - let compressed_bit_width = u64::from(u64::BITS - max_value.leading_zeros()).max(1); + let compressed_bit_width = u64::from(u64::BITS - max_value.leading_zeros()); fixed_width.block_info.0.write().unwrap().insert( Stat::BitWidth, std::sync::Arc::new(UInt64Array::from(vec![compressed_bit_width])), @@ -765,6 +765,30 @@ mod test { ); } + #[test] + fn test_inline_block_bitpacking_preserves_zero_bit_width() { + let codec = InlineBitpacking::new(8); + let input = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_u8; ELEMS_PER_CHUNK as usize]), + bits_per_value: 8, + num_values: ELEMS_PER_CHUNK, + block_info: BlockInfo::new(), + }); + + let payload = BlockCompressor::compress(&codec, input).unwrap().unwrap(); + assert_eq!(payload.as_ref(), &[0]); + + let decoded = + BlockDecompressor::decompress(&codec, Some(payload), ELEMS_PER_CHUNK).unwrap(); + let DataBlock::FixedWidth(decoded) = decoded else { + panic!("Expected FixedWidth block"); + }; + assert_eq!( + decoded.data.borrow_to_typed_view::().as_ref(), + vec![0; ELEMS_PER_CHUNK as usize] + ); + } + // 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.