From c4ca75edf18ff601e2e00dec19c53a330c6cce0a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 5 Aug 2026 21:44:59 +0800 Subject: [PATCH] feat(encoding): compose generic block codecs --- rust/lance-encoding/src/compression.rs | 413 +++++++++++++--- .../src/encodings/physical/block.rs | 17 + .../src/encodings/physical/constant.rs | 127 ++++- .../src/encodings/physical/general.rs | 75 ++- .../src/encodings/physical/rle.rs | 448 ++++++++++++++++++ 5 files changed, 1007 insertions(+), 73 deletions(-) diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 50f679b0947..657f5ca9582 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -42,13 +42,17 @@ use crate::{ byte_stream_split::{ ByteStreamSplitDecompressor, ByteStreamSplitEncoder, should_use_bss, }, - constant::ConstantDecompressor, + constant::{ConstantBlockDecompressor, ConstantDecompressor}, delta::DeltaDecompressor, + dictionary::BlockDictionaryDecompressor, fsst::{ FsstMiniBlockDecompressor, FsstMiniBlockEncoder, FsstPerValueDecompressor, FsstPerValueEncoder, }, - general::{GeneralMiniBlockCompressor, GeneralMiniBlockDecompressor}, + general::{ + FixedWidthGeneralBlockDecompressor, GeneralMiniBlockCompressor, + GeneralMiniBlockDecompressor, + }, packed::{ PackedStructFixedWidthMiniBlockDecompressor, PackedStructFixedWidthMiniBlockEncoder, PackedStructVariablePerValueDecompressor, @@ -57,8 +61,8 @@ use crate::{ }, range::RangeDecompressor, rle::{ - RleChildDecompressor, RleDecompressor, RleEncoder, RunLengthWidth, - rle_encoded_size, select_run_length_width, + BlockRleDecompressor, BlockRleRunCount, RleChildDecompressor, RleDecompressor, + RleEncoder, RunLengthWidth, rle_encoded_size, select_run_length_width, }, value::{ValueDecompressor, ValueEncoder}, }, @@ -1195,81 +1199,283 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Ok(Box::new(general_decompressor)) } Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)), - Compression::Range(range) => { - if !matches!(range.uncompressed_bits_per_value, 32 | 64) { - return Err(Error::invalid_input(format!( - "Range only supports 32 or 64-bit values, got {}", - range.uncompressed_bits_per_value - ))); - } - if range.uncompressed_bits_per_value == 32 && range.start > u32::MAX as u64 { - return Err(Error::invalid_input(format!( - "Range start {} exceeds u32::MAX", - range.start - ))); - } - if range.step == 0 { - return Err(Error::invalid_input("Range step must be positive")); + Compression::Range(range) => create_fixed_width_block_decompressor( + description, + range.uncompressed_bits_per_value, + ), + Compression::Delta(delta) => create_fixed_width_block_decompressor( + description, + delta.uncompressed_bits_per_value, + ), + _ => todo!(), + } + } +} + +/// Builds the bounded decoder tree used by generic unsigned block sequences. +/// +/// This is intentionally separate from [`DecompressionStrategy::create_block_decompressor`]: +/// stable 2.1/2.2 block readers keep their existing grammar and framing, while +/// the 2.3 offset container opts into this typed grammar explicitly. +pub(crate) fn create_fixed_width_block_decompressor( + description: &CompressiveEncoding, + expected_bits_per_value: u64, +) -> Result> { + if !matches!(expected_bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Generic block decoding only supports 32 or 64-bit values, got {expected_bits_per_value}" + ))); + } + create_fixed_width_block_decompressor_inner(description, expected_bits_per_value, true) +} + +fn create_fixed_width_block_decompressor_inner( + description: &CompressiveEncoding, + expected_bits_per_value: u64, + allow_composite: bool, +) -> Result> { + let compression = description + .compression + .as_ref() + .ok_or_else(|| Error::invalid_input("Block encoding is missing its compression variant"))?; + match compression { + Compression::Flat(flat) => { + if flat.bits_per_value != expected_bits_per_value { + return Err(Error::invalid_input(format!( + "Flat declares {}-bit values, expected {expected_bits_per_value}", + flat.bits_per_value + ))); + } + if flat.data.is_some() { + return Err(Error::invalid_input( + "Generic block Flat cannot contain buffer compression", + )); + } + Ok(Box::new(ValueDecompressor::from_flat(flat))) + } + Compression::Constant(constant) => { + let value = decode_fixed_width_constant(constant, expected_bits_per_value)?; + Ok(Box::new(ConstantBlockDecompressor::new( + expected_bits_per_value, + value, + ))) + } + Compression::Range(range) => { + if range.uncompressed_bits_per_value != expected_bits_per_value { + return Err(Error::invalid_input(format!( + "Range declares {}-bit values, expected {expected_bits_per_value}", + range.uncompressed_bits_per_value + ))); + } + Ok(Box::new(RangeDecompressor::new( + expected_bits_per_value, + range.start, + range.step, + ))) + } + Compression::InlineBitpacking(bitpacking) => { + if bitpacking.uncompressed_bits_per_value != expected_bits_per_value { + return Err(Error::invalid_input(format!( + "Inline bitpacking declares {}-bit values, expected {expected_bits_per_value}", + bitpacking.uncompressed_bits_per_value + ))); + } + #[cfg(feature = "bitpacking")] + { + Ok(Box::new(InlineBitpacking::from_description(bitpacking))) + } + #[cfg(not(feature = "bitpacking"))] + { + Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )) + } + } + Compression::OutOfLineBitpacking(bitpacking) => { + if bitpacking.uncompressed_bits_per_value != expected_bits_per_value { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking declares {}-bit values, expected {expected_bits_per_value}", + bitpacking.uncompressed_bits_per_value + ))); + } + let values = bitpacking.values.as_deref().ok_or_else(|| { + Error::invalid_input("Out-of-line bitpacking is missing its values encoding") + })?; + let compressed_bits = match values.compression.as_ref() { + Some(Compression::Flat(flat)) if flat.data.is_none() => flat.bits_per_value, + _ => { + return Err(Error::invalid_input( + "Out-of-line bitpacking values must use plain Flat encoding", + )); } - Ok(Box::new(RangeDecompressor::new( - range.uncompressed_bits_per_value, - range.start, - range.step, + }; + if compressed_bits == 0 || compressed_bits >= expected_bits_per_value { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking width {compressed_bits} is invalid for {expected_bits_per_value}-bit values" + ))); + } + #[cfg(feature = "bitpacking")] + { + Ok(Box::new(OutOfLineBitpacking::new( + compressed_bits, + expected_bits_per_value, ))) } - Compression::Delta(delta) => { - let bits_per_value = delta.uncompressed_bits_per_value; - if !matches!(bits_per_value, 32 | 64) { - return Err(Error::invalid_input(format!( - "Delta only supports 32 or 64-bit values, got {bits_per_value}" - ))); - } - if bits_per_value == 32 && delta.base > u32::MAX as u64 { - return Err(Error::invalid_input(format!( - "Delta base {} exceeds u32::MAX", - delta.base - ))); + #[cfg(not(feature = "bitpacking"))] + { + Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )) + } + } + Compression::Delta(delta) if allow_composite => { + if delta.uncompressed_bits_per_value != expected_bits_per_value { + return Err(Error::invalid_input(format!( + "Delta declares {}-bit values, expected {expected_bits_per_value}", + delta.uncompressed_bits_per_value + ))); + } + let child = delta.deltas.as_deref().ok_or_else(|| { + Error::invalid_input("Delta is missing its deltas child encoding") + })?; + let child = + create_fixed_width_block_decompressor_inner(child, expected_bits_per_value, false)?; + Ok(Box::new(DeltaDecompressor::new( + expected_bits_per_value, + delta.base, + child, + ))) + } + Compression::General(general) if allow_composite => { + let child = general.values.as_deref().ok_or_else(|| { + Error::invalid_input("General block compression is missing its values encoding") + })?; + if !matches!( + child.compression.as_ref(), + Some(Compression::Flat(flat)) + if flat.bits_per_value == expected_bits_per_value && flat.data.is_none() + ) { + return Err(Error::invalid_input( + "Generic General compression must wrap plain Flat values", + )); + } + let child = + create_fixed_width_block_decompressor_inner(child, expected_bits_per_value, false)?; + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input("General block compression is missing its compression config") + })?; + let scheme = compression.scheme().try_into()?; + let config = CompressionConfig::new(scheme, compression.level); + Ok(Box::new(FixedWidthGeneralBlockDecompressor::new( + child, + config, + expected_bits_per_value, + ))) + } + Compression::Rle(rle) if allow_composite => { + let values = rle.values.as_deref().ok_or_else(|| { + Error::invalid_input("RLE compression is missing its values encoding") + })?; + let run_lengths = rle.run_lengths.as_deref().ok_or_else(|| { + Error::invalid_input("RLE compression is missing its run-length encoding") + })?; + let run_count = match run_lengths.compression.as_ref() { + Some(Compression::Constant(constant)) => { + BlockRleRunCount::ConstantRunLength(decode_fixed_width_constant(constant, 32)?) } - let child = delta.deltas.as_deref().ok_or_else(|| { - Error::invalid_input("Delta is missing its deltas child encoding") - })?; - let child_bits = match child.compression.as_ref() { - Some(Compression::Flat(flat)) => flat.bits_per_value, - Some(Compression::Range(range)) => range.uncompressed_bits_per_value, - Some(Compression::InlineBitpacking(bitpacking)) => { - bitpacking.uncompressed_bits_per_value - } - Some(Compression::OutOfLineBitpacking(bitpacking)) => { - bitpacking.uncompressed_bits_per_value + Some(Compression::Range(range)) if range.uncompressed_bits_per_value == 32 => { + BlockRleRunCount::RangeRunLengths { + start: range.start, + step: range.step, } - Some(other) => { - return Err(Error::invalid_input(format!( - "Delta does not support a {} child", - compression_name(other) - ))); - } - None => { - return Err(Error::invalid_input( - "Delta child is missing its compression variant", - )); - } - }; - if child_bits != bits_per_value { - return Err(Error::invalid_input(format!( - "Delta child declares {child_bits}-bit values, expected {bits_per_value}" - ))); } - let child = self.create_block_decompressor(child)?; - Ok(Box::new(DeltaDecompressor::new( - bits_per_value, - delta.base, - child, - ))) - } - _ => todo!(), + _ if matches!(values.compression.as_ref(), Some(Compression::Flat(flat)) if flat.bits_per_value == expected_bits_per_value && flat.data.is_none()) => { + BlockRleRunCount::FlatValues + } + _ if matches!(run_lengths.compression.as_ref(), Some(Compression::Flat(flat)) if flat.bits_per_value == 32 && flat.data.is_none()) => { + BlockRleRunCount::FlatRunLengths + } + _ => { + return Err(Error::invalid_input( + "RLE requires metadata run lengths or one plain Flat child to recover the run count", + )); + } + }; + let values = create_fixed_width_block_decompressor_inner( + values, + expected_bits_per_value, + false, + )?; + let run_lengths = create_fixed_width_block_decompressor_inner(run_lengths, 32, false)?; + Ok(Box::new(BlockRleDecompressor::try_new( + expected_bits_per_value, + values, + run_lengths, + run_count, + )?)) + } + Compression::Dictionary(dictionary) if allow_composite => { + let indices = dictionary.indices.as_deref().ok_or_else(|| { + Error::invalid_input("Dictionary is missing its indices encoding") + })?; + let items = dictionary + .items + .as_deref() + .ok_or_else(|| Error::invalid_input("Dictionary is missing its items encoding"))?; + let indices = create_fixed_width_block_decompressor_inner(indices, 32, false)?; + let items = + create_fixed_width_block_decompressor_inner(items, expected_bits_per_value, false)?; + Ok(Box::new(BlockDictionaryDecompressor::try_new( + expected_bits_per_value, + dictionary.num_dictionary_items, + indices, + items, + )?)) } + other => Err(Error::invalid_input(format!( + "{} is not allowed at this position in a generic block encoding", + compression_name(other) + ))), } } + +fn decode_fixed_width_constant( + constant: &crate::format::pb21::Constant, + bits_per_value: u64, +) -> Result { + let value = constant + .value + .as_ref() + .ok_or_else(|| Error::invalid_input("Typed Constant is missing its scalar value"))?; + let expected_bytes = usize::try_from(bits_per_value / 8) + .map_err(|_| Error::invalid_input("Constant scalar width does not fit usize"))?; + if value.len() != expected_bytes { + return Err(Error::invalid_input(format!( + "Constant scalar has {} bytes, expected {expected_bytes}", + value.len() + ))); + } + Ok(match bits_per_value { + 32 => u64::from(u32::from_le_bytes( + value + .as_ref() + .try_into() + .expect("constant width was checked"), + )), + 64 => u64::from_le_bytes( + value + .as_ref() + .try_into() + .expect("constant width was checked"), + ), + _ => { + return Err(Error::invalid_input(format!( + "Constant only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + }) +} + pub(crate) fn create_rle_decompressor( rle: &crate::format::pb21::Rle, decompression_strategy: &dyn DecompressionStrategy, @@ -2785,4 +2991,71 @@ mod tests { "RLE should not be used for V2.1" ); } + + #[test] + fn generic_block_rle_round_trips_metadata_children() { + use crate::encodings::physical::{ + constant::ConstantEncoder, range::RangeEncoder, rle::BlockRleEncoder, + }; + + let values = vec![10_u64, 10, 20, 20, 30, 30]; + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(values.clone()), + num_values: values.len() as u64, + block_info: BlockInfo::default(), + }); + let encoder = BlockRleEncoder::try_new( + 64, + Box::new(RangeEncoder::new(64, 10, 10)), + Box::new(ConstantEncoder::new(32, 2)), + ) + .unwrap(); + let (payload, encoding) = encoder.compress(block).unwrap(); + assert!(payload.is_none()); + + let decoder = create_fixed_width_block_decompressor(&encoding, 64).unwrap(); + let decoded = decoder.decompress(None, values.len() as u64).unwrap(); + assert_eq!( + decoded + .as_fixed_width() + .unwrap() + .data + .borrow_to_typed_slice::() + .as_ref(), + values + ); + } + + #[test] + fn generic_block_dictionary_uses_the_typed_decoder_tree() { + use crate::encodings::physical::{dictionary::BlockDictionaryEncoder, value::ValueEncoder}; + + let values = vec![10_u64, 20, 10, 20]; + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(values.clone()), + num_values: values.len() as u64, + block_info: BlockInfo::default(), + }); + let encoder = BlockDictionaryEncoder::try_new( + 64, + Arc::from([10, 20]), + Box::new(ValueEncoder::default()), + Box::new(ValueEncoder::default()), + ) + .unwrap(); + let (payload, encoding) = encoder.compress(block).unwrap(); + let decoder = create_fixed_width_block_decompressor(&encoding, 64).unwrap(); + let decoded = decoder.decompress(payload, values.len() as u64).unwrap(); + assert_eq!( + decoded + .as_fixed_width() + .unwrap() + .data + .borrow_to_typed_slice::() + .as_ref(), + values + ); + } } diff --git a/rust/lance-encoding/src/encodings/physical/block.rs b/rust/lance-encoding/src/encodings/physical/block.rs index a645988a147..9e2eb94abd7 100644 --- a/rust/lance-encoding/src/encodings/physical/block.rs +++ b/rust/lance-encoding/src/encodings/physical/block.rs @@ -141,6 +141,23 @@ 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<()>; + + fn decompress_exact( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + expected_bytes: usize, + ) -> Result<()> { + self.decompress(input_buf, output_buf)?; + if output_buf.len() != expected_bytes { + return Err(Error::invalid_input(format!( + "General compression produced {} bytes, expected {expected_bytes}", + output_buf.len() + ))); + } + Ok(()) + } + fn config(&self) -> CompressionConfig; } diff --git a/rust/lance-encoding/src/encodings/physical/constant.rs b/rust/lance-encoding/src/encodings/physical/constant.rs index dd153f789d5..c71b200434c 100644 --- a/rust/lance-encoding/src/encodings/physical/constant.rs +++ b/rust/lance-encoding/src/encodings/physical/constant.rs @@ -5,11 +5,134 @@ use crate::{ buffer::LanceBuffer, - compression::{BlockDecompressor, FixedPerValueDecompressor, require_no_block_payload}, + compression::{ + BlockCompressor, BlockDecompressor, FixedPerValueDecompressor, require_no_block_payload, + }, data::{AllNullDataBlock, ConstantDataBlock, DataBlock, FixedWidthDataBlock}, + encodings::physical::{checked_fixed_values, try_vec_with_capacity}, + format::{ProtobufUtils21, pb21::CompressiveEncoding}, }; -use lance_core::Result; +use lance_core::{Error, Result}; + +/// Metadata-only compressor for a repeated unsigned `u32` or `u64` value. +#[derive(Debug)] +pub struct ConstantEncoder { + bits_per_value: u64, + value: u64, +} + +impl ConstantEncoder { + pub fn new(bits_per_value: u64, value: u64) -> Self { + Self { + bits_per_value, + value, + } + } +} + +impl BlockCompressor for ConstantEncoder { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Constant block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.bits_per_value { + return Err(Error::invalid_input(format!( + "Constant codec expects {}-bit values, got {}", + self.bits_per_value, data.bits_per_value + ))); + } + let scalar = match self.bits_per_value { + 32 => { + let value = u32::try_from(self.value).map_err(|_| { + Error::invalid_input(format!("Constant value {} exceeds u32::MAX", self.value)) + })?; + if checked_fixed_values::(&data, "Constant input")? + .iter() + .any(|candidate| *candidate != value) + { + return Err(Error::invalid_input( + "Constant input contains a different value", + )); + } + bytes::Bytes::copy_from_slice(&value.to_le_bytes()) + } + 64 => { + if checked_fixed_values::(&data, "Constant input")? + .iter() + .any(|candidate| *candidate != self.value) + { + return Err(Error::invalid_input( + "Constant input contains a different value", + )); + } + bytes::Bytes::copy_from_slice(&self.value.to_le_bytes()) + } + bits_per_value => { + return Err(Error::invalid_input(format!( + "Constant block compression only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + }; + Ok((None, ProtobufUtils21::constant(Some(scalar)))) + } +} + +/// Materializes a metadata-only constant as a typed fixed-width block. +#[derive(Debug)] +pub(crate) struct ConstantBlockDecompressor { + bits_per_value: u64, + value: u64, +} + +impl ConstantBlockDecompressor { + pub(crate) fn new(bits_per_value: u64, value: u64) -> Self { + Self { + bits_per_value, + value, + } + } +} + +impl BlockDecompressor for ConstantBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + require_no_block_payload(data, "Constant")?; + let output_len = usize::try_from(num_values) + .map_err(|_| Error::invalid_input("Constant output cardinality does not fit usize"))?; + let data = match self.bits_per_value { + 32 => { + let value = u32::try_from(self.value).map_err(|_| { + Error::invalid_input(format!("Constant value {} exceeds u32::MAX", self.value)) + })?; + let mut values = try_vec_with_capacity::(num_values, "Constant output")?; + values.resize(output_len, value); + LanceBuffer::reinterpret_vec(values) + } + 64 => { + let mut values = try_vec_with_capacity::(num_values, "Constant output")?; + values.resize(output_len, self.value); + LanceBuffer::reinterpret_vec(values) + } + bits_per_value => { + return Err(Error::invalid_input(format!( + "Constant block decompression only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + }; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.bits_per_value, + data, + num_values, + block_info: Default::default(), + })) + } + + fn requires_payload(&self) -> bool { + false + } +} /// A decompressor for constant-encoded data #[derive(Debug)] diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index bea6a85cb01..c65f93dcfdd 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -6,7 +6,9 @@ use log::trace; use crate::{ Result, buffer::LanceBuffer, - compression::MiniBlockDecompressor, + compression::{ + BlockCompressor, BlockDecompressor, MiniBlockDecompressor, require_block_payload, + }, data::DataBlock, encodings::{ logical::primitive::miniblock::{ @@ -16,6 +18,77 @@ use crate::{ }, format::{ProtobufUtils21, pb21::CompressiveEncoding}, }; +use lance_core::Error; + +/// General-purpose compressor that wraps one payload-bearing block codec. +#[derive(Debug)] +pub struct GeneralBlockCompressor { + child: Box, + compression: CompressionConfig, +} + +impl GeneralBlockCompressor { + pub fn new(child: Box, compression: CompressionConfig) -> Self { + Self { child, compression } + } +} + +impl BlockCompressor for GeneralBlockCompressor { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let (payload, child_encoding) = self.child.compress(data)?; + let payload = payload.ok_or_else(|| { + Error::invalid_input("General block compression requires a payload-bearing child") + })?; + let compressor = GeneralBufferCompressor::get_compressor(self.compression)?; + let mut compressed = Vec::new(); + compressor.compress(&payload, &mut compressed)?; + Ok(( + Some(LanceBuffer::from(compressed)), + ProtobufUtils21::wrapped(self.compression, child_encoding)?, + )) + } +} + +/// Bounded fixed-width decoder for a general-compressed block payload. +#[derive(Debug)] +pub(crate) struct FixedWidthGeneralBlockDecompressor { + child: Box, + compression: CompressionConfig, + bits_per_value: u64, +} + +impl FixedWidthGeneralBlockDecompressor { + pub(crate) fn new( + child: Box, + compression: CompressionConfig, + bits_per_value: u64, + ) -> Self { + Self { + child, + compression, + bits_per_value, + } + } +} + +impl BlockDecompressor for FixedWidthGeneralBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "General fixed-width block")?; + let expected_bytes = usize::try_from(num_values) + .ok() + .and_then(|num_values| { + num_values.checked_mul(usize::try_from(self.bits_per_value / 8).ok()?) + }) + .ok_or_else(|| { + Error::invalid_input("General fixed-width output length overflows usize") + })?; + let compressor = GeneralBufferCompressor::get_compressor(self.compression)?; + let mut decompressed = Vec::new(); + compressor.decompress_exact(&data, &mut decompressed, expected_bytes)?; + self.child + .decompress(Some(LanceBuffer::from(decompressed)), num_values) + } +} /// A miniblock compressor that wraps another miniblock compressor and applies /// general-purpose compression (LZ4, Zstd) to the resulting buffers. diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index 0215e392584..35bb9c3e000 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -1857,6 +1857,454 @@ impl RleDecompressor { } } +// Generic block RLE uses the same outer frame as the stable block codec while +// allowing each child to be any concrete block codec. +pub(crate) const GENERIC_BLOCK_RLE_HEADER_BYTES: usize = 8; + +/// Block RLE compressor with concrete value and run-length children. +#[derive(Debug)] +pub struct BlockRleEncoder { + bits_per_value: u64, + values: Box, + run_lengths: Box, +} + +impl BlockRleEncoder { + pub fn try_new( + bits_per_value: u64, + values: Box, + run_lengths: Box, + ) -> Result { + if !matches!(bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Generic block RLE only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + Ok(Self { + bits_per_value, + values, + run_lengths, + }) + } +} + +impl BlockCompressor for BlockRleEncoder { + fn compress(&self, data: DataBlock) -> Result<(Option, CompressiveEncoding)> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Generic block RLE requires fixed-width data", + )); + }; + let (run_values, run_lengths) = materialize_generic_runs(&data, self.bits_per_value)?; + let (values_payload, values_encoding) = + self.values.compress(DataBlock::FixedWidth(run_values))?; + let (lengths_payload, lengths_encoding) = self + .run_lengths + .compress(DataBlock::FixedWidth(run_lengths))?; + let encoding = ProtobufUtils21::rle(values_encoding, lengths_encoding); + if values_payload.is_none() && lengths_payload.is_none() { + return Ok((None, encoding)); + } + + let values_payload = values_payload.unwrap_or_else(LanceBuffer::empty); + let lengths_payload = lengths_payload.unwrap_or_else(LanceBuffer::empty); + let capacity = GENERIC_BLOCK_RLE_HEADER_BYTES + .checked_add(values_payload.len()) + .and_then(|capacity| capacity.checked_add(lengths_payload.len())) + .ok_or_else(|| { + Error::invalid_input("Generic block RLE frame length overflows usize") + })?; + let mut output = Vec::new(); + output.try_reserve_exact(capacity).map_err(|error| { + Error::invalid_input(format!( + "Generic block RLE could not reserve {capacity} frame bytes: {error}" + )) + })?; + 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)), encoding)) + } +} + +fn materialize_generic_runs( + data: &FixedWidthDataBlock, + bits_per_value: u64, +) -> Result<(FixedWidthDataBlock, FixedWidthDataBlock)> { + if data.bits_per_value != bits_per_value { + return Err(Error::invalid_input(format!( + "Generic block RLE expects {bits_per_value}-bit values, got {}", + data.bits_per_value + ))); + } + if data.num_values == 0 { + return Err(Error::invalid_input( + "Generic block RLE cannot encode an empty sequence", + )); + } + + fn collect(values: &[T]) -> (Vec, Vec) { + let mut run_values = Vec::new(); + let mut run_lengths = Vec::new(); + let mut current = values[0]; + let mut length = 1_u32; + for value in values.iter().copied().skip(1) { + if value == current && length < u32::MAX { + length += 1; + } else { + run_values.push(current); + run_lengths.push(length); + current = value; + length = 1; + } + } + run_values.push(current); + run_lengths.push(length); + (run_values, run_lengths) + } + + let (values, lengths, num_runs) = match bits_per_value { + 32 => { + let input = crate::encodings::physical::checked_fixed_values::( + data, + "Generic block RLE input", + )?; + let (values, lengths) = collect(&input); + let num_runs = values.len() as u64; + (LanceBuffer::reinterpret_vec(values), lengths, num_runs) + } + 64 => { + let input = crate::encodings::physical::checked_fixed_values::( + data, + "Generic block RLE input", + )?; + let (values, lengths) = collect(&input); + let num_runs = values.len() as u64; + (LanceBuffer::reinterpret_vec(values), lengths, num_runs) + } + _ => unreachable!("generic block RLE width was validated at construction"), + }; + Ok(( + FixedWidthDataBlock { + bits_per_value, + data: values, + num_values: num_runs, + block_info: BlockInfo::default(), + }, + FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(lengths), + num_values: num_runs, + block_info: BlockInfo::default(), + }, + )) +} + +/// Descriptor-derived source used to recover the RLE child cardinality. +#[derive(Debug, Clone, Copy)] +pub(crate) enum BlockRleRunCount { + ConstantRunLength(u64), + RangeRunLengths { start: u64, step: u64 }, + FlatValues, + FlatRunLengths, +} + +/// Block RLE decompressor with concrete value and run-length children. +#[derive(Debug)] +pub(crate) struct BlockRleDecompressor { + bits_per_value: u64, + values: Box, + run_lengths: Box, + run_count: BlockRleRunCount, +} + +impl BlockRleDecompressor { + pub(crate) fn try_new( + bits_per_value: u64, + values: Box, + run_lengths: Box, + run_count: BlockRleRunCount, + ) -> Result { + if !matches!(bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Generic block RLE only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + Ok(Self { + bits_per_value, + values, + run_lengths, + run_count, + }) + } +} + +impl BlockDecompressor for BlockRleDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + if num_values == 0 { + return Err(Error::invalid_input( + "Generic block RLE cannot decode an empty sequence", + )); + } + let values_have_payload = self.values.requires_payload(); + let lengths_have_payload = self.run_lengths.requires_payload(); + let (values_payload, lengths_payload) = + split_generic_rle_payload(data, values_have_payload, lengths_have_payload)?; + let run_count = infer_generic_run_count( + self.run_count, + values_payload.as_ref(), + lengths_payload.as_ref(), + self.bits_per_value, + num_values, + )?; + let values = self.values.decompress(values_payload, run_count)?; + let lengths = self.run_lengths.decompress(lengths_payload, run_count)?; + expand_generic_runs(values, lengths, self.bits_per_value, num_values, run_count) + } + + fn requires_payload(&self) -> bool { + self.values.requires_payload() || self.run_lengths.requires_payload() + } +} + +fn split_generic_rle_payload( + data: Option, + values_have_payload: bool, + lengths_have_payload: bool, +) -> Result<(Option, Option)> { + if !values_have_payload && !lengths_have_payload { + if data.is_some() { + return Err(Error::invalid_input( + "Metadata-only generic block RLE expects no payload", + )); + } + return Ok((None, None)); + } + let data = require_block_payload(data, "Generic block RLE")?; + if data.len() < GENERIC_BLOCK_RLE_HEADER_BYTES { + return Err(Error::invalid_input(format!( + "Generic block RLE payload has {} bytes, shorter than its {GENERIC_BLOCK_RLE_HEADER_BYTES}-byte header", + data.len() + ))); + } + let values_size = u64::from_le_bytes( + data[..GENERIC_BLOCK_RLE_HEADER_BYTES] + .try_into() + .expect("generic block RLE header length was checked"), + ); + let values_size = usize::try_from(values_size).map_err(|_| { + Error::invalid_input("Generic block RLE values payload length does not fit usize") + })?; + let lengths_start = GENERIC_BLOCK_RLE_HEADER_BYTES + .checked_add(values_size) + .ok_or_else(|| Error::invalid_input("Generic block RLE values payload end overflows"))?; + if lengths_start > data.len() { + return Err(Error::invalid_input(format!( + "Generic block RLE values payload ends at {lengths_start}, beyond {} bytes", + data.len() + ))); + } + if !values_have_payload && values_size != 0 { + return Err(Error::invalid_input(format!( + "Metadata-only RLE values child has {values_size} framed payload bytes" + ))); + } + if !lengths_have_payload && lengths_start != data.len() { + return Err(Error::invalid_input(format!( + "Metadata-only RLE run-length child has {} framed payload bytes", + data.len() - lengths_start + ))); + } + Ok(( + values_have_payload + .then(|| data.slice_with_length(GENERIC_BLOCK_RLE_HEADER_BYTES, values_size)), + lengths_have_payload + .then(|| data.slice_with_length(lengths_start, data.len() - lengths_start)), + )) +} + +fn infer_generic_run_count( + source: BlockRleRunCount, + values_payload: Option<&LanceBuffer>, + lengths_payload: Option<&LanceBuffer>, + bits_per_value: u64, + num_values: u64, +) -> Result { + let run_count = match source { + BlockRleRunCount::ConstantRunLength(length) => { + if length == 0 || !num_values.is_multiple_of(length) { + return Err(Error::invalid_input(format!( + "RLE constant run length {length} does not divide {num_values} values" + ))); + } + num_values / length + } + BlockRleRunCount::RangeRunLengths { start, step } => { + if start == 0 || step == 0 { + return Err(Error::invalid_input( + "RLE range run lengths must be positive", + )); + } + let target = u128::from(num_values); + let mut low = 1_u64; + let mut high = num_values; + let mut found = None; + while low <= high { + let count = low + (high - low) / 2; + let count128 = u128::from(count); + let sum = count128 + .checked_mul( + u128::from(start) + .checked_mul(2) + .and_then(|start| { + u128::from(count - 1) + .checked_mul(u128::from(step)) + .and_then(|tail| start.checked_add(tail)) + }) + .unwrap_or(u128::MAX), + ) + .map(|sum| sum / 2) + .unwrap_or(u128::MAX); + match sum.cmp(&target) { + std::cmp::Ordering::Less => low = count + 1, + std::cmp::Ordering::Greater => high = count - 1, + std::cmp::Ordering::Equal => { + found = Some(count); + break; + } + } + } + found.ok_or_else(|| { + Error::invalid_input(format!( + "RLE range run lengths start={start} step={step} do not sum to {num_values}" + )) + })? + } + BlockRleRunCount::FlatValues => flat_generic_run_count( + values_payload.ok_or_else(|| { + Error::invalid_input("RLE values payload is required to infer the run count") + })?, + bits_per_value, + num_values, + "values", + )?, + BlockRleRunCount::FlatRunLengths => flat_generic_run_count( + lengths_payload.ok_or_else(|| { + Error::invalid_input("RLE run-length payload is required to infer the run count") + })?, + 32, + num_values, + "run lengths", + )?, + }; + if run_count == 0 { + return Err(Error::invalid_input( + "Generic block RLE describes zero runs", + )); + } + Ok(run_count) +} + +fn flat_generic_run_count( + payload: &LanceBuffer, + bits_per_value: u64, + max_runs: u64, + label: &str, +) -> Result { + let bytes_per_value = usize::try_from(bits_per_value / 8) + .map_err(|_| Error::invalid_input("RLE child width does not fit usize"))?; + if bytes_per_value == 0 || !payload.len().is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input(format!( + "RLE {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_runs { + return Err(Error::invalid_input(format!( + "RLE {label} payload contains {run_count} runs, exceeding {max_runs}" + ))); + } + Ok(run_count) +} + +fn expand_generic_runs( + values: DataBlock, + lengths: DataBlock, + bits_per_value: u64, + 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(lengths) = lengths else { + return Err(Error::invalid_input( + "RLE run lengths decoded to a non fixed-width block", + )); + }; + if values.bits_per_value != bits_per_value + || values.num_values != run_count + || lengths.bits_per_value != 32 + || lengths.num_values != run_count + { + return Err(Error::invalid_input( + "RLE child cardinality or bit width does not match its descriptor", + )); + } + let lengths = + crate::encodings::physical::checked_fixed_values::(&lengths, "RLE run lengths")?; + let total = lengths + .iter() + .enumerate() + .try_fold(0_u64, |total, (index, length)| { + if *length == 0 { + return Err(Error::invalid_input(format!( + "RLE run length at index {index} is zero" + ))); + } + total.checked_add(u64::from(*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 output = match bits_per_value { + 32 => { + let values = + crate::encodings::physical::checked_fixed_values::(&values, "RLE values")?; + let mut output = + crate::encodings::physical::try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in values.iter().zip(lengths.iter()) { + output.extend(std::iter::repeat_n(*value, *length as usize)); + } + LanceBuffer::reinterpret_vec(output) + } + 64 => { + let values = + crate::encodings::physical::checked_fixed_values::(&values, "RLE values")?; + let mut output = + crate::encodings::physical::try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in values.iter().zip(lengths.iter()) { + output.extend(std::iter::repeat_n(*value, *length as usize)); + } + LanceBuffer::reinterpret_vec(output) + } + _ => unreachable!("generic block RLE width was validated at construction"), + }; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value, + data: output, + num_values, + block_info: BlockInfo::default(), + })) +} + #[cfg(test)] mod tests { use std::sync::Arc;