diff --git a/protos/table.proto b/protos/table.proto index d298809d5d8..318f8566aae 100644 --- a/protos/table.proto +++ b/protos/table.proto @@ -109,10 +109,12 @@ message Manifest { // should not attempt to read the dataset. // // Known flags: - // * 1: deletion files are present - // * 2: row ids are stable and stored as part of the fragment metadata. - // * 4: use v2 format (deprecated) - // * 8: table config is present + // * 1 << 0: deletion files are present + // * 1 << 1: row ids are stable and stored as part of the fragment metadata. + // * 1 << 2: use v2 format (deprecated) + // * 1 << 3: table config is present + // * 1 << 4: dataset uses multiple base paths + // * 1 << 5: transaction file writes are disabled uint64 reader_feature_flags = 9; // Feature flags for writers. diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index e5d617e44ef..79e32daa40d 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -51,7 +51,10 @@ use crate::{ PackedStructVariablePerValueEncoder, VariablePackedStructFieldDecoder, VariablePackedStructFieldKind, }, - rle::{RleDecompressor, RleEncoder}, + rle::{ + RleDecompressor, RleEncoder, RunLengthWidth, select_run_length_width, + select_run_length_width_and_size, + }, value::{ValueDecompressor, ValueEncoder}, }, }, @@ -165,6 +168,7 @@ fn try_bss_for_mini_block( fn try_rle_for_mini_block( data: &FixedWidthDataBlock, params: &CompressionFieldParams, + use_rle_v2: bool, ) -> Option> { let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { @@ -188,16 +192,22 @@ fn try_rle_for_mini_block( return None; } - // Estimate the encoded size. - // - // RLE stores (value, run_length) pairs. Run lengths are u8 and long runs are split into - // multiple entries of up to 255 values. We don't know the run length distribution here, - // so we conservatively account for splitting with an upper bound. let num_values = data.num_values; - let estimated_pairs = (run_count.saturating_add(num_values / 255)).min(num_values); - let raw_bytes = (num_values as u128) * (type_size as u128); - let rle_bytes = (estimated_pairs as u128) * ((type_size + 1) as u128); + let (run_length_width, rle_bytes) = if use_rle_v2 { + select_run_length_width_and_size(&data.data, data.num_values, data.bits_per_value).ok()? + } else { + // Estimate the encoded size. + // + // Compatibility RLE stores (value, u8 run_length) pairs. Long runs are split into + // multiple entries of up to 255 values. We don't know the run length distribution here, + // so we conservatively account for splitting with an upper bound. + let estimated_pairs = (run_count.saturating_add(num_values / 255)).min(num_values); + ( + RunLengthWidth::U8, + (estimated_pairs as u128) * ((type_size + 1) as u128), + ) + }; if rle_bytes < raw_bytes { #[cfg(feature = "bitpacking")] @@ -208,7 +218,9 @@ fn try_rle_for_mini_block( return None; } } - return Some(Box::new(RleEncoder::new())); + return Some(Box::new(RleEncoder::with_run_length_width( + run_length_width, + ))); } None } @@ -217,14 +229,15 @@ fn try_rle_for_block( data: &FixedWidthDataBlock, version: LanceFileVersion, params: &CompressionFieldParams, -) -> Option<(Box, CompressiveEncoding)> { + use_rle_v2: bool, +) -> Result, CompressiveEncoding)>> { if version < LanceFileVersion::V2_2 { - return None; + return Ok(None); } let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { - return None; + return Ok(None); } let run_count = data.expect_single_stat::(Stat::RunCount); @@ -233,14 +246,19 @@ fn try_rle_for_block( .unwrap_or(DEFAULT_RLE_COMPRESSION_THRESHOLD); if (run_count as f64) < (data.num_values as f64) * threshold { - let compressor = Box::new(RleEncoder::new()); + let run_length_width = if use_rle_v2 { + select_run_length_width(&data.data, data.num_values, data.bits_per_value)? + } else { + RunLengthWidth::U8 + }; + let compressor = Box::new(RleEncoder::with_run_length_width(run_length_width)); let encoding = ProtobufUtils21::rle( ProtobufUtils21::flat(bits, None), - ProtobufUtils21::flat(/*bits_per_value=*/ 8, None), + ProtobufUtils21::flat(run_length_width.bits_per_value(), None), ); - return Some((compressor, encoding)); + return Ok(Some((compressor, encoding))); } - None + Ok(None) } fn try_bitpack_for_mini_block(_data: &FixedWidthDataBlock) -> Option> { @@ -393,6 +411,10 @@ impl DefaultCompressionStrategy { self } + fn use_rle_v2(&self) -> bool { + self.version.resolve() >= LanceFileVersion::V2_3 + } + /// Parse compression parameters from field metadata fn parse_field_metadata(field: &Field, version: &LanceFileVersion) -> CompressionFieldParams { let mut params = CompressionFieldParams::default(); @@ -456,7 +478,7 @@ impl DefaultCompressionStrategy { } let base = try_bss_for_mini_block(data, params) - .or_else(|| try_rle_for_mini_block(data, params)) + .or_else(|| try_rle_for_mini_block(data, params, self.use_rle_v2())) .or_else(|| try_bitpack_for_mini_block(data)) .unwrap_or_else(|| Box::new(ValueEncoder::default())); @@ -664,7 +686,7 @@ impl CompressionStrategy for DefaultCompressionStrategy { match data { DataBlock::FixedWidth(fixed_width) => { if let Some((compressor, encoding)) = - try_rle_for_block(fixed_width, self.version, &field_params) + try_rle_for_block(fixed_width, self.version, &field_params, self.use_rle_v2())? { return Ok((compressor, encoding)); } @@ -815,8 +837,11 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Ok(Box::new(ValueDecompressor::from_fsl(fsl))) } Compression::Rle(rle) => { - let bits_per_value = validate_rle_compression(rle)?; - Ok(Box::new(RleDecompressor::new(bits_per_value))) + let (bits_per_value, run_length_width) = validate_rle_compression(rle)?; + Ok(Box::new(RleDecompressor::with_run_length_width( + bits_per_value, + run_length_width, + ))) } Compression::ByteStreamSplit(bss) => { let Compression::Flat(values) = @@ -1005,15 +1030,18 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Ok(Box::new(general_decompressor)) } Compression::Rle(rle) => { - let bits_per_value = validate_rle_compression(rle)?; - Ok(Box::new(RleDecompressor::new(bits_per_value))) + let (bits_per_value, run_length_width) = validate_rle_compression(rle)?; + Ok(Box::new(RleDecompressor::with_run_length_width( + bits_per_value, + run_length_width, + ))) } _ => todo!(), } } } -/// Validates RLE compression format and extracts bits_per_value -fn validate_rle_compression(rle: &crate::format::pb21::Rle) -> Result { +/// Validates RLE compression format and extracts value and run length widths. +fn validate_rle_compression(rle: &crate::format::pb21::Rle) -> Result<(u64, RunLengthWidth)> { let values = rle .values .as_ref() @@ -1043,14 +1071,22 @@ fn validate_rle_compression(rle: &crate::format::pb21::Rle) -> Result { )); }; - if run_lengths.bits_per_value != 8 { + if !matches!(values.bits_per_value, 8 | 16 | 32 | 64) { return Err(Error::invalid_input(format!( - "RLE compression only supports 8-bit run lengths, got {}", - run_lengths.bits_per_value + "RLE compression only supports 8, 16, 32, or 64-bit values, got {}", + values.bits_per_value ))); } - Ok(values.bits_per_value) + let run_length_width = + RunLengthWidth::from_bits(run_lengths.bits_per_value).ok_or_else(|| { + Error::invalid_input(format!( + "RLE compression only supports 8, 16, or 32-bit run lengths, got {}", + run_lengths.bits_per_value + )) + })?; + + Ok((values.bits_per_value, run_length_width)) } #[cfg(test)] @@ -1659,6 +1695,83 @@ mod tests { assert!(debug_str.contains("RleEncoder")); } + #[test] + fn test_rle_v2_miniblock_selects_u16_run_lengths() { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![7i32; 1000]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 1000, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + 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 Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { + panic!("expected RLE encoding"); + }; + let Compression::Flat(run_lengths) = rle + .run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("expected flat run lengths"); + }; + assert_eq!(run_lengths.bits_per_value, 16); + } + + #[test] + fn test_rle_v2_uses_selected_width_cost_before_bitpacking() { + let mut metadata = HashMap::new(); + metadata.insert(RLE_THRESHOLD_META_KEY.to_string(), "1.0".to_string()); + metadata.insert(BSS_META_KEY.to_string(), "off".to_string()); + let mut field = create_test_field("test_column", DataType::Int32); + field.metadata = metadata; + + let values = vec![0i32; 4096]; + let mut data = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(values), + num_values: 4096, + block_info: BlockInfo::default(), + }; + data.compute_stat(); + let data = DataBlock::FixedWidth(data); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let debug_str = format!("{compressor:?}"); + assert!(debug_str.contains("RleEncoder")); + + let (_compressed, encoding) = compressor.compress(data).unwrap(); + let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { + panic!("expected RLE encoding"); + }; + let Compression::Flat(run_lengths) = rle + .run_lengths + .as_ref() + .unwrap() + .compression + .as_ref() + .unwrap() + else { + panic!("expected flat run lengths"); + }; + assert_eq!(run_lengths.bits_per_value, 16); + } + #[test] fn test_field_metadata_override_params() { // Set up params with one configuration diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index 4d58f72e71a..53c61928870 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -161,7 +161,7 @@ mod tests { // Small data with RLE - should not compress due to size threshold TestCase { name: "small_rle_data", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -173,7 +173,7 @@ mod tests { // Large repeated data with RLE + LZ4 TestCase { name: "large_rle_lz4", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -185,7 +185,7 @@ mod tests { // Large repeated data with RLE + Zstd TestCase { name: "large_rle_zstd", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Zstd, level: Some(3), @@ -403,7 +403,7 @@ mod tests { // Test that small buffers don't get compressed let small_test = TestCase { name: "small_buffer_no_compression", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -496,7 +496,7 @@ mod tests { // RLE produces 2 buffers (values and lengths), test that both are handled correctly let data = create_repeated_i32_block(vec![1; 100]); let compressor = GeneralMiniBlockCompressor::new( - Box::new(RleEncoder), + Box::new(RleEncoder::new()), CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -519,7 +519,7 @@ mod tests { // Test case 1: 32-bit RLE data let test_32 = TestCase { name: "rle_32bit_with_general_wrapper", - inner_encoder: Box::new(RleEncoder), + inner_encoder: Box::new(RleEncoder::new()), compression: CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -532,7 +532,7 @@ mod tests { // For 32-bit RLE, the compression strategy should automatically wrap it // Let's directly test the compressor let compressor = GeneralMiniBlockCompressor::new( - Box::new(RleEncoder), + Box::new(RleEncoder::new()), CompressionConfig { scheme: CompressionScheme::Lz4, level: None, @@ -589,7 +589,7 @@ mod tests { let block_64 = DataBlock::from_array(array_64); let compressor_64 = GeneralMiniBlockCompressor::new( - Box::new(RleEncoder), + Box::new(RleEncoder::new()), CompressionConfig { scheme: CompressionScheme::Lz4, level: None, diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index 88e27bf954f..66960fb3dc4 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -10,7 +10,7 @@ //! RLE uses a dual-buffer format to store compressed data: //! //! - **Values Buffer**: Stores unique values in their original data type -//! - **Lengths Buffer**: Stores the repeat count for each value as u8 +//! - **Lengths Buffer**: Stores the repeat count for each value as u8, u16, or u32 //! //! ### Example //! @@ -18,13 +18,13 @@ //! //! Encoded as: //! - Values buffer: `[1, 2, 3]` (3 × 4 bytes for i32) -//! - Lengths buffer: `[3, 2, 4]` (3 × 1 byte for u8) +//! - Lengths buffer: `[3, 2, 4]` (3 × 1 byte for u8 in compatibility mode) //! //! ### Long Run Handling //! -//! When a run exceeds 255 values, it is split into multiple runs of 255 -//! followed by a final run with the remainder. For example, a run of 1000 -//! identical values becomes 4 runs: [255, 255, 255, 235]. +//! In compatibility mode, when a run exceeds 255 values, it is split into multiple +//! runs of 255 followed by a final run with the remainder. RLE v2 can use u16 or +//! u32 run lengths to reduce this splitting. //! //! ## Supported Types //! @@ -70,13 +70,195 @@ use crate::format::pb21::CompressiveEncoding; use lance_core::{Error, Result}; +/// Width used to encode RLE run lengths. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RunLengthWidth { + /// Compatibility mode. Runs longer than 255 values are split. + U8, + /// RLE v2 mode for runs up to 65,535 values per entry. + U16, + /// RLE v2 mode for runs up to 4,294,967,295 values per entry. + U32, +} + +impl RunLengthWidth { + pub(crate) fn from_bits(bits_per_value: u64) -> Option { + match bits_per_value { + 8 => Some(Self::U8), + 16 => Some(Self::U16), + 32 => Some(Self::U32), + _ => None, + } + } + + pub(crate) fn bits_per_value(self) -> u64 { + match self { + Self::U8 => 8, + Self::U16 => 16, + Self::U32 => 32, + } + } + + fn bytes_per_value(self) -> usize { + match self { + Self::U8 => 1, + Self::U16 => 2, + Self::U32 => 4, + } + } + + fn max_run_length(self) -> u64 { + match self { + Self::U8 => u8::MAX as u64, + Self::U16 => u16::MAX as u64, + Self::U32 => u32::MAX as u64, + } + } + + fn write_length(self, length: u64, dst: &mut Vec) { + match self { + Self::U8 => dst.push(length as u8), + Self::U16 => dst.extend_from_slice(&(length as u16).to_le_bytes()), + Self::U32 => dst.extend_from_slice(&(length as u32).to_le_bytes()), + } + } + + fn read_length(self, bytes: &[u8]) -> u64 { + match self { + Self::U8 => bytes[0] as u64, + Self::U16 => u16::from_le_bytes([bytes[0], bytes[1]]) as u64, + Self::U32 => u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as u64, + } + } +} + +/// Select the lowest-cost run length width for a fixed-width data block. +pub(crate) fn select_run_length_width( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, +) -> Result { + select_run_length_width_and_size(data, num_values, bits_per_value).map(|(width, _)| width) +} + +/// Select the lowest-cost run length width and return the encoded size estimate. +pub(crate) fn select_run_length_width_and_size( + data: &LanceBuffer, + num_values: u64, + bits_per_value: u64, +) -> Result<(RunLengthWidth, u128)> { + match bits_per_value { + 8 => select_run_length_width_generic::(data, num_values), + 16 => select_run_length_width_generic::(data, num_values), + 32 => select_run_length_width_generic::(data, num_values), + 64 => select_run_length_width_generic::(data, num_values), + _ => Err(Error::invalid_input_source( + format!("RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}") + .into(), + )), + } +} + +fn select_run_length_width_generic( + data: &LanceBuffer, + num_values: u64, +) -> Result<(RunLengthWidth, u128)> +where + T: bytemuck::Pod + PartialEq + Copy + ArrowNativeType, +{ + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + if num_values == 0 { + return Ok((RunLengthWidth::U8, 0)); + } + + let type_size = std::mem::size_of::(); + let expected_bytes = num_values.checked_mul(type_size).ok_or_else(|| { + Error::invalid_input_source( + format!("RLE input byte length overflow: {num_values} values of {type_size} bytes") + .into(), + ) + })?; + if data.len() != expected_bytes { + return Err(Error::invalid_input_source( + format!( + "RLE input data size mismatch: {} bytes for {} values of {} bytes", + data.len(), + num_values, + type_size + ) + .into(), + )); + } + + let values_ref = data.borrow_to_typed_slice::(); + let values: &[T] = values_ref.as_ref(); + let mut costs = [0_u128; 3]; + + let mut current_value = values[0]; + let mut current_length = 1_u64; + for &value in values.iter().skip(1) { + if value == current_value { + current_length += 1; + } else { + accumulate_width_costs(current_length, type_size, &mut costs); + current_value = value; + current_length = 1; + } + } + accumulate_width_costs(current_length, type_size, &mut costs); + + let widths = [RunLengthWidth::U8, RunLengthWidth::U16, RunLengthWidth::U32]; + let mut best_idx = 0usize; + let mut best_cost = costs[0]; + for (idx, &cost) in costs.iter().enumerate().skip(1) { + if cost < best_cost { + best_idx = idx; + best_cost = cost; + } + } + Ok((widths[best_idx], best_cost)) +} + +fn accumulate_width_costs(run_length: u64, type_size: usize, costs: &mut [u128; 3]) { + let widths = [RunLengthWidth::U8, RunLengthWidth::U16, RunLengthWidth::U32]; + // The current encoder uses miniblock-sized chunks for both miniblock and block paths. + let max_segment_values = *MAX_MINIBLOCK_VALUES; + let mut remaining = run_length; + while remaining > 0 { + let segment = remaining.min(max_segment_values); + for (idx, width) in widths.iter().enumerate() { + let entries = segment.div_ceil(width.max_run_length()); + costs[idx] += (entries as u128) * ((type_size + width.bytes_per_value()) as u128); + } + remaining -= segment; + } +} + /// RLE encoder for miniblock format -#[derive(Debug, Default)] -pub struct RleEncoder; +#[derive(Debug)] +pub struct RleEncoder { + run_length_width: RunLengthWidth, +} + +impl Default for RleEncoder { + fn default() -> Self { + Self::new() + } +} impl RleEncoder { pub fn new() -> Self { - Self + Self { + run_length_width: RunLengthWidth::U8, + } + } + + pub(crate) fn with_run_length_width(run_length_width: RunLengthWidth) -> Self { + Self { run_length_width } } fn encode_data( @@ -89,17 +271,23 @@ impl RleEncoder { return Ok((Vec::new(), Vec::new())); } + let num_values = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; let bytes_per_value = (bits_per_value / 8) as usize; + let bytes_per_length = self.run_length_width.bytes_per_value(); // Pre-allocate global buffers with estimated capacity // Assume average compression ratio of ~10:1 (10 values per run) - let estimated_runs = num_values as usize / 10; + let estimated_runs = num_values / 10; let mut all_values = Vec::with_capacity(estimated_runs * bytes_per_value); - let mut all_lengths = Vec::with_capacity(estimated_runs); + let mut all_lengths = Vec::with_capacity(estimated_runs * bytes_per_length); let mut chunks = Vec::new(); let mut offset = 0usize; - let mut values_remaining = num_values as usize; + let mut values_remaining = num_values; while values_remaining > 0 { let values_start = all_values.len(); @@ -134,7 +322,14 @@ impl RleEncoder { &mut all_values, &mut all_lengths, ), - _ => unreachable!("RLE encoding bits_per_value must be 8, 16, 32 or 64"), + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE encoding bits_per_value must be 8, 16, 32, or 64, got {bits_per_value}" + ) + .into(), + )); + } }; if values_processed == 0 { @@ -175,22 +370,7 @@ impl RleEncoder { )) } - /// Encodes a chunk of data using RLE compression with dynamic boundary detection. - /// - /// This function processes values sequentially, detecting runs (sequences of identical values) - /// and encoding them as (value, length) pairs. It dynamically determines whether this chunk - /// should be the last chunk based on how many values were processed. - /// - /// # Key Features: - /// - Tracks byte usage to ensure we don't exceed MAX_MINIBLOCK_BYTES - /// - Maintains power-of-2 checkpoints for non-last chunks - /// - Splits long runs (>255) into multiple entries - /// - Dynamically determines if this is the last chunk - /// - /// # Returns: - /// - num_runs: Number of runs encoded - /// - values_processed: Number of input values processed - /// - is_last_chunk: Whether this chunk processed all remaining values + /// Encodes the largest valid mini-block prefix from `offset`. fn encode_chunk_rolling( &self, data: &LanceBuffer, @@ -203,7 +383,6 @@ impl RleEncoder { T: bytemuck::Pod + PartialEq + Copy + std::fmt::Debug + ArrowNativeType, { let type_size = std::mem::size_of::(); - let chunk_start = offset * type_size; let max_by_count = *MAX_MINIBLOCK_VALUES as usize; let max_values = values_remaining.min(max_by_count); @@ -217,112 +396,101 @@ impl RleEncoder { let chunk_buffer = data.slice_with_length(chunk_start, chunk_len); let typed_data_ref = chunk_buffer.borrow_to_typed_slice::(); let typed_data: &[T] = typed_data_ref.as_ref(); + let max_values = max_values.min(typed_data.len()); if typed_data.is_empty() { return (0, 0, false); } - // Record starting positions for this chunk let values_start = all_values.len(); + let all_remaining_values_fit = values_remaining <= max_by_count; + let encoded_size = self.encoded_size(&typed_data[..max_values]); + let (values_to_encode, is_last_chunk) = if all_remaining_values_fit + && encoded_size <= MAX_MINIBLOCK_BYTES as usize + { + (max_values, true) + } else if let Some(values_to_encode) = self.largest_power_of_two_prefix::(typed_data) { + (values_to_encode, false) + } else { + return (0, 0, false); + }; - let mut current_value = typed_data[0]; - let mut current_length = 1u64; - let mut bytes_used = 0usize; - let mut total_values_encoded = 0usize; // Track total encoded values + self.encode_values(&typed_data[..values_to_encode], all_values, all_lengths); - // Power-of-2 checkpoints for ensuring non-last chunks have valid sizes. - // - // We start from a slightly larger minimum checkpoint for smaller types since - // they encode more compactly and are less likely to hit MAX_MINIBLOCK_BYTES. - let min_checkpoint_log2 = match type_size { - 1 => 8, // 256 - 2 => 7, // 128 - _ => 6, // 64 - }; - let max_checkpoint_log2 = (values_remaining.min(*MAX_MINIBLOCK_VALUES as usize)) - .next_power_of_two() - .ilog2(); - let mut checkpoint_log2 = min_checkpoint_log2; + let num_runs = (all_values.len() - values_start) / type_size; + (num_runs, values_to_encode, is_last_chunk) + } + + fn largest_power_of_two_prefix(&self, values: &[T]) -> Option + where + T: bytemuck::Pod + PartialEq + Copy, + { + let max_prefix = values.len().min(*MAX_MINIBLOCK_VALUES as usize); + let mut prefix = 1usize << max_prefix.ilog2(); + while prefix > 1 { + if self.encoded_size(&values[..prefix]) <= MAX_MINIBLOCK_BYTES as usize { + return Some(prefix); + } + prefix >>= 1; + } + None + } + + fn encoded_size(&self, values: &[T]) -> usize + where + T: bytemuck::Pod + PartialEq + Copy, + { + if values.is_empty() { + return 0; + } - // Save state at checkpoints so we can roll back if needed - let mut last_checkpoint_state = None; + let mut current_value = values[0]; + let mut current_length = 1u64; + let mut encoded_size = 0usize; - for &value in typed_data[1..].iter() { + for &value in values.iter().skip(1) { if value == current_value { current_length += 1; } else { - // Calculate space needed (may need multiple u8s if run > 255) - let run_chunks = current_length.div_ceil(255) as usize; - let bytes_needed = run_chunks * (type_size + 1); - - // Stop if adding this run would exceed byte limit - if bytes_used + bytes_needed > MAX_MINIBLOCK_BYTES as usize { - if let Some((val_pos, len_pos, _, checkpoint_values)) = last_checkpoint_state { - // Roll back to last power-of-2 checkpoint - all_values.truncate(val_pos); - all_lengths.truncate(len_pos); - let num_runs = (val_pos - values_start) / type_size; - return (num_runs, checkpoint_values, false); - } - break; - } - - bytes_used += self.add_run(¤t_value, current_length, all_values, all_lengths); - total_values_encoded += current_length as usize; + encoded_size += self.run_size::(current_length); current_value = value; current_length = 1; } - - // Check if we reached a power-of-2 checkpoint. - while checkpoint_log2 <= max_checkpoint_log2 { - let checkpoint_values = 1usize << checkpoint_log2; - if checkpoint_values > values_remaining || total_values_encoded < checkpoint_values - { - break; - } - last_checkpoint_state = Some(( - all_values.len(), - all_lengths.len(), - bytes_used, - checkpoint_values, - )); - checkpoint_log2 += 1; - } } + encoded_size += self.run_size::(current_length); + encoded_size + } - // After the loop, we always have a pending run that needs to be added - // unless we've exceeded the byte limit - if current_length > 0 { - let run_chunks = current_length.div_ceil(255) as usize; - let bytes_needed = run_chunks * (type_size + 1); + fn run_size(&self, length: u64) -> usize + where + T: bytemuck::Pod, + { + let type_size = std::mem::size_of::(); + let run_chunks = length.div_ceil(self.run_length_width.max_run_length()) as usize; + run_chunks * (type_size + self.run_length_width.bytes_per_value()) + } - if bytes_used + bytes_needed <= MAX_MINIBLOCK_BYTES as usize { - let _ = self.add_run(¤t_value, current_length, all_values, all_lengths); - total_values_encoded += current_length as usize; - } + fn encode_values(&self, values: &[T], all_values: &mut Vec, all_lengths: &mut Vec) + where + T: bytemuck::Pod + PartialEq + Copy, + { + if values.is_empty() { + return; } - // Determine if we've processed all remaining values - let is_last_chunk = total_values_encoded == values_remaining; - - // Non-last chunks must have power-of-2 values for miniblock format - if !is_last_chunk { - if total_values_encoded.is_power_of_two() { - // Already at power-of-2 boundary - } else if let Some((val_pos, len_pos, _, checkpoint_values)) = last_checkpoint_state { - // Roll back to last valid checkpoint - all_values.truncate(val_pos); - all_lengths.truncate(len_pos); - let num_runs = (val_pos - values_start) / type_size; - return (num_runs, checkpoint_values, false); + let mut current_value = values[0]; + let mut current_length = 1u64; + + for &value in values.iter().skip(1) { + if value == current_value { + current_length += 1; } else { - // No valid checkpoint, can't create a valid chunk - return (0, 0, false); + self.add_run(¤t_value, current_length, all_values, all_lengths); + current_value = value; + current_length = 1; } } - - let num_runs = (all_values.len() - values_start) / type_size; - (num_runs, total_values_encoded, is_last_chunk) + self.add_run(¤t_value, current_length, all_values, all_lengths); } fn add_run( @@ -337,24 +505,26 @@ impl RleEncoder { { let value_bytes = bytemuck::bytes_of(value); let type_size = std::mem::size_of::(); - let num_full_chunks = (length / 255) as usize; - let remainder = (length % 255) as u8; + let max_run_length = self.run_length_width.max_run_length(); + let num_full_chunks = (length / max_run_length) as usize; + let remainder = length % max_run_length; let total_chunks = num_full_chunks + if remainder > 0 { 1 } else { 0 }; all_values.reserve(total_chunks * type_size); - all_lengths.reserve(total_chunks); + all_lengths.reserve(total_chunks * self.run_length_width.bytes_per_value()); for _ in 0..num_full_chunks { all_values.extend_from_slice(value_bytes); - all_lengths.push(255); + self.run_length_width + .write_length(max_run_length, all_lengths); } if remainder > 0 { all_values.extend_from_slice(value_bytes); - all_lengths.push(remainder); + self.run_length_width.write_length(remainder, all_lengths); } - total_chunks * (type_size + 1) + total_chunks * (type_size + self.run_length_width.bytes_per_value()) } } @@ -376,7 +546,7 @@ impl MiniBlockCompressor for RleEncoder { let encoding = ProtobufUtils21::rle( ProtobufUtils21::flat(bits_per_value, None), - ProtobufUtils21::flat(/*bits_per_value=*/ 8, None), + ProtobufUtils21::flat(self.run_length_width.bits_per_value(), None), ); Ok((compressed, encoding)) @@ -418,11 +588,25 @@ impl BlockCompressor for RleEncoder { #[derive(Debug)] pub struct RleDecompressor { bits_per_value: u64, + run_length_width: RunLengthWidth, } impl RleDecompressor { pub fn new(bits_per_value: u64) -> Self { - Self { bits_per_value } + Self { + bits_per_value, + run_length_width: RunLengthWidth::U8, + } + } + + pub(crate) fn with_run_length_width( + bits_per_value: u64, + run_length_width: RunLengthWidth, + ) -> Self { + Self { + bits_per_value, + run_length_width, + } } fn decode_data(&self, data: Vec, num_values: u64) -> Result { @@ -453,7 +637,15 @@ impl RleDecompressor { 16 => self.decode_generic::(values_buffer, lengths_buffer, num_values)?, 32 => self.decode_generic::(values_buffer, lengths_buffer, num_values)?, 64 => self.decode_generic::(values_buffer, lengths_buffer, num_values)?, - _ => unreachable!("RLE decoding bits_per_value must be 8, 16, 32, 64, or 128"), + _ => { + return Err(Error::invalid_input_source( + format!( + "RLE decoding bits_per_value must be 8, 16, 32, or 64, got {}", + self.bits_per_value + ) + .into(), + )); + } }; Ok(DataBlock::FixedWidth(FixedWidthDataBlock { @@ -474,6 +666,7 @@ impl RleDecompressor { T: bytemuck::Pod + Copy + std::fmt::Debug + ArrowNativeType, { let type_size = std::mem::size_of::(); + let length_size = self.run_length_width.bytes_per_value(); if values_buffer.is_empty() || lengths_buffer.is_empty() { if num_values == 0 { @@ -485,19 +678,22 @@ impl RleDecompressor { } } - if !values_buffer.len().is_multiple_of(type_size) || lengths_buffer.is_empty() { + if !values_buffer.len().is_multiple_of(type_size) + || !lengths_buffer.len().is_multiple_of(length_size) + { return Err(Error::invalid_input_source(format!( - "Invalid buffer sizes for RLE {} decoding: values {} bytes (not divisible by {}), lengths {} bytes", + "Invalid buffer sizes for RLE {} decoding: values {} bytes (not divisible by {}), lengths {} bytes (not divisible by {})", std::any::type_name::(), values_buffer.len(), type_size, - lengths_buffer.len() + lengths_buffer.len(), + length_size ) .into())); } let num_runs = values_buffer.len() / type_size; - let num_length_entries = lengths_buffer.len(); + let num_length_entries = lengths_buffer.len() / length_size; if num_runs != num_length_entries { return Err(Error::invalid_input_source( format!( @@ -510,39 +706,56 @@ impl RleDecompressor { let values_ref = values_buffer.borrow_to_typed_slice::(); let values: &[T] = values_ref.as_ref(); - let lengths: &[u8] = lengths_buffer.as_ref(); - - let expected_value_count = num_values as usize; - let mut decoded: Vec = Vec::with_capacity(expected_value_count); - - for (value, &length) in values.iter().zip(lengths.iter()) { - if decoded.len() == expected_value_count { - break; - } + let lengths = lengths_buffer.as_ref(); + let expected_value_count = usize::try_from(num_values).map_err(|_| { + Error::invalid_input_source( + format!("RLE num_values does not fit in usize: {num_values}").into(), + ) + })?; + let mut decoded_value_count = 0usize; + for length_bytes in lengths.chunks_exact(length_size) { + let length = self.run_length_width.read_length(length_bytes); if length == 0 { return Err(Error::invalid_input_source( "RLE decoding encountered a zero run length".into(), )); } - - let remaining = expected_value_count - decoded.len(); - let write_len = (length as usize).min(remaining); - - decoded.resize(decoded.len() + write_len, *value); + let length = usize::try_from(length).map_err(|_| { + Error::invalid_input_source( + format!("RLE run length does not fit in usize: {length}").into(), + ) + })?; + decoded_value_count = decoded_value_count.checked_add(length).ok_or_else(|| { + Error::invalid_input_source("RLE run length sum overflowed usize".into()) + })?; + if decoded_value_count > expected_value_count { + return Err(Error::invalid_input_source( + format!( + "RLE decoding overflowed expected value count: produced at least {}, expected {}", + decoded_value_count, expected_value_count + ) + .into(), + )); + } } - if decoded.len() != expected_value_count { + if decoded_value_count != expected_value_count { return Err(Error::invalid_input_source( format!( "RLE decoding produced {} values, expected {}", - decoded.len(), - expected_value_count + decoded_value_count, expected_value_count ) .into(), )); } + let mut decoded: Vec = Vec::with_capacity(expected_value_count); + for (value, length_bytes) in values.iter().zip(lengths.chunks_exact(length_size)) { + let length = self.run_length_width.read_length(length_bytes) as usize; + decoded.resize(decoded.len() + length, *value); + } + trace!( "RLE decoded {} {} values", num_values, @@ -641,6 +854,58 @@ mod tests { assert_eq!(lengths_buffer.len(), 6); } + #[test] + fn test_rle_v2_u16_miniblock_encoding() { + let encoder = RleEncoder::with_run_length_width(RunLengthWidth::U16); + + let data = vec![42i32; 1000]; + let array = Int32Array::from(data); + let (compressed, encoding) = + MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + + assert_eq!(compressed.data[0].len(), 4); + assert_eq!(compressed.data[1].len(), 2); + assert_eq!(compressed.data[1].as_ref(), &1000u16.to_le_bytes()); + + let rle = match encoding.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle, + other => panic!("expected RLE encoding, got {other:?}"), + }; + let run_lengths = rle.run_lengths.as_ref().unwrap(); + let flat = match run_lengths.compression.as_ref().unwrap() { + crate::format::pb21::compressive_encoding::Compression::Flat(flat) => flat, + other => panic!("expected flat run lengths, got {other:?}"), + }; + assert_eq!(flat.bits_per_value, 16); + + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let decompressed = MiniBlockDecompressor::decompress( + &decompressor, + compressed.data, + compressed.num_values, + ) + .unwrap(); + match decompressed { + DataBlock::FixedWidth(block) => { + let values = block.data.borrow_to_typed_slice::(); + assert_eq!(values.as_ref(), vec![42i32; 1000]); + } + _ => panic!("Expected FixedWidth block"), + } + } + + #[test] + fn test_select_run_length_width_prefers_u16_for_long_runs() { + let data = vec![7i32; 300]; + let bytes = data + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let width = + select_run_length_width(&LanceBuffer::from(bytes), data.len() as u64, 32).unwrap(); + assert_eq!(width, RunLengthWidth::U16); + } + // ========== Round-trip Tests for Different Types ========== #[test] @@ -760,6 +1025,61 @@ mod tests { ); } + #[test] + fn test_u16_length_buffer_must_be_aligned() { + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let values = LanceBuffer::from(vec![1, 0, 0, 0]); + let lengths = LanceBuffer::from(vec![5]); + let result = MiniBlockDecompressor::decompress(&decompressor, vec![values, lengths], 5); + assert!(matches!(&result, Err(Error::InvalidInput { .. }))); + assert!( + result + .unwrap_err() + .to_string() + .contains("not divisible by 2") + ); + } + + #[test] + fn test_rle_rejects_underflow_overflow_and_zero_lengths() { + let decompressor = RleDecompressor::with_run_length_width(32, RunLengthWidth::U16); + let value = LanceBuffer::from(1i32.to_le_bytes().to_vec()); + + let underflow = MiniBlockDecompressor::decompress( + &decompressor, + vec![ + value.clone(), + LanceBuffer::from(4u16.to_le_bytes().to_vec()), + ], + 5, + ) + .unwrap_err(); + assert!(underflow.to_string().contains("produced 4 values")); + + let overflow = MiniBlockDecompressor::decompress( + &decompressor, + vec![ + value.clone(), + LanceBuffer::from(6u16.to_le_bytes().to_vec()), + ], + 5, + ) + .unwrap_err(); + assert!( + overflow + .to_string() + .contains("overflowed expected value count") + ); + + let zero = MiniBlockDecompressor::decompress( + &decompressor, + vec![value, LanceBuffer::from(0u16.to_le_bytes().to_vec())], + 5, + ) + .unwrap_err(); + assert!(zero.to_string().contains("zero run length")); + } + #[test] fn test_empty_data_handling() { let encoder = RleEncoder::new(); diff --git a/rust/lance/src/dataset/tests/dataset_io.rs b/rust/lance/src/dataset/tests/dataset_io.rs index c0c90fc71c9..7e52f572e9b 100644 --- a/rust/lance/src/dataset/tests/dataset_io.rs +++ b/rust/lance/src/dataset/tests/dataset_io.rs @@ -15,7 +15,7 @@ use crate::session::Session; use crate::{Dataset, Error, Result}; use lance_table::format::DataStorageFormat; -use crate::dataset::write::{WriteMode, WriteParams}; +use crate::dataset::write::{CommitBuilder, InsertBuilder, WriteMode, WriteParams}; use arrow::array::as_struct_array; use arrow::compute::concat_batches; use arrow_array::RecordBatch; @@ -890,6 +890,164 @@ async fn test_write_manifest( assert!(matches!(write_result, Err(Error::NotSupported { .. }))); } +#[tokio::test] +async fn test_rle_v2_v23_write_and_append() { + let test_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![7; 1000]))], + ) + .unwrap(); + + let batches = RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema.clone()); + let mut dataset = Dataset::write( + batches, + &test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }), + ) + .await + .unwrap(); + + let manifest = read_manifest( + dataset.object_store.as_ref(), + &dataset + .commit_handler + .resolve_latest_location(&dataset.base, dataset.object_store.as_ref()) + .await + .unwrap() + .path, + None, + ) + .await + .unwrap(); + assert_eq!( + manifest.data_storage_format.lance_file_version().unwrap(), + LanceFileVersion::V2_3 + ); + + let append_batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![9; 1000]))], + ) + .unwrap(); + let append_batches = + RecordBatchIterator::new(vec![Ok(append_batch)].into_iter(), schema.clone()); + dataset = Dataset::write( + append_batches, + &test_uri, + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await + .unwrap(); + + assert_eq!( + dataset + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::V2_3 + ); + + let actual = dataset.scan().try_into_batch().await.unwrap(); + let expected = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from( + [vec![7; 1000], vec![9; 1000]].concat(), + ))], + ) + .unwrap(); + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn test_rle_v2_uncommitted_create_commits_v23_storage() { + let test_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![7; 1000]))], + ) + .unwrap(); + + let transaction = InsertBuilder::new(test_uri.as_str()) + .with_params(&WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }) + .execute_uncommitted(vec![batch]) + .await + .unwrap(); + + let dataset = CommitBuilder::new(test_uri.as_str()) + .execute(transaction) + .await + .unwrap(); + assert_eq!( + dataset + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::V2_3 + ); +} + +#[tokio::test] +async fn test_rle_v2_shallow_clone_preserves_v23_storage() { + let test_uri = TempStrDir::default(); + let clone_uri = TempStrDir::default(); + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "i", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![7; 1000]))], + ) + .unwrap(); + + let mut dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch)].into_iter(), schema), + &test_uri, + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + ..Default::default() + }), + ) + .await + .unwrap(); + + let clone = dataset + .shallow_clone(clone_uri.as_str(), dataset.version().version, None) + .await + .unwrap(); + assert_eq!( + clone + .manifest + .data_storage_format + .lance_file_version() + .unwrap(), + LanceFileVersion::V2_3 + ); +} + #[rstest] #[tokio::test] async fn append_dataset( diff --git a/rust/lance/src/dataset/write.rs b/rust/lance/src/dataset/write.rs index ff0a119158c..65efe1737a9 100644 --- a/rust/lance/src/dataset/write.rs +++ b/rust/lance/src/dataset/write.rs @@ -313,6 +313,7 @@ pub struct WriteParams { /// /// Newer versions are more efficient but the data can only be read by more recent versions /// of lance. + /// Lance file version 2.3 enables RLE v2 run length widths by default. /// /// If not specified then the latest stable version will be used. pub data_storage_version: Option, diff --git a/rust/lance/src/dataset/write/commit.rs b/rust/lance/src/dataset/write/commit.rs index baad71b3e39..8a4d131af37 100644 --- a/rust/lance/src/dataset/write/commit.rs +++ b/rust/lance/src/dataset/write/commit.rs @@ -345,7 +345,6 @@ impl<'a> CommitBuilder<'a> { } else { self.use_stable_row_ids.unwrap_or(false) }; - // Validate storage format matches existing dataset if let Some(ds) = dest.dataset() && let Some(storage_format) = self.storage_format @@ -503,7 +502,6 @@ impl<'a> CommitBuilder<'a> { }, read_version, tag: None, - //TODO: handle batch transaction merges in the future transaction_properties: None, }; let dataset = self.execute(merged.clone()).await?;