From e7e21a290717686a4a18f583e4b8ca46c2b73735 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Tue, 28 Jul 2026 17:18:13 +0800 Subject: [PATCH] feat(encoding): encode generic mini-block offsets --- rust/lance-encoding/src/compression/block.rs | 18 +- .../src/compression/block/factory.rs | 2 - .../src/compression/block/fixed.rs | 2 - .../encodings/logical/primitive/miniblock.rs | 28 + .../src/encodings/physical/binary.rs | 1203 +++++++++++++++-- .../src/encodings/physical/constant.rs | 8 +- .../src/encodings/physical/delta.rs | 1 - .../src/encodings/physical/dictionary.rs | 22 +- .../src/encodings/physical/general.rs | 9 +- .../src/encodings/physical/range.rs | 8 +- .../src/encodings/physical/rle.rs | 12 +- .../src/encodings/physical/value.rs | 3 - 12 files changed, 1132 insertions(+), 184 deletions(-) diff --git a/rust/lance-encoding/src/compression/block.rs b/rust/lance-encoding/src/compression/block.rs index 99904e5d9b5..8974bc866d6 100644 --- a/rust/lance-encoding/src/compression/block.rs +++ b/rust/lance-encoding/src/compression/block.rs @@ -65,20 +65,16 @@ impl BlockValueType { mod factory; pub(crate) mod fixed; -#[cfg(test)] -pub(crate) use factory::encode_scalar; -#[cfg(all(test, feature = "bitpacking"))] -pub(crate) use factory::out_of_line_payload_bytes; pub(crate) use factory::{ - create_block_decompressor, infer_block_value_type, validate_fixed_payload_len, + create_block_decompressor, encode_scalar, infer_block_value_type, 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}; +pub(crate) use factory::{ + out_of_line_payload_bytes, validate_inline_bitpacking_payload, validate_out_of_line_payload, +}; +pub(crate) use fixed::{ + fixed_block, fixed_from_u64_values, read_unsigned_values, visit_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 index b389a4d9ca8..2c03646a3a0 100644 --- a/rust/lance-encoding/src/compression/block/factory.rs +++ b/rust/lance-encoding/src/compression/block/factory.rs @@ -3,7 +3,6 @@ //! Fallible construction of concrete generic block decompressors. -#[cfg(test)] use bytes::Bytes; use super::*; @@ -525,7 +524,6 @@ fn decode_scalar(bytes: &[u8], value_type: BlockValueType, label: &str) -> Resul }) } -#[cfg(test)] pub fn encode_scalar(value: u64, value_type: BlockValueType) -> Result { if value > value_type.max_value() { return Err(Error::invalid_input(format!( diff --git a/rust/lance-encoding/src/compression/block/fixed.rs b/rust/lance-encoding/src/compression/block/fixed.rs index d5288d571ef..409bd5a0c3d 100644 --- a/rust/lance-encoding/src/compression/block/fixed.rs +++ b/rust/lance-encoding/src/compression/block/fixed.rs @@ -16,7 +16,6 @@ pub fn fixed_block(value_type: BlockValueType, num_values: u64, data: LanceBuffe }) } -#[cfg(test)] pub fn fixed_from_u64_values( values: &[u64], value_type: BlockValueType, @@ -63,7 +62,6 @@ pub fn fixed_from_u64_values( }) } -#[cfg(any(test, feature = "bitpacking"))] pub fn visit_unsigned_values( block: &FixedWidthDataBlock, value_type: BlockValueType, diff --git a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs index 9461229310b..1d25d1fdfb0 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs @@ -75,6 +75,19 @@ impl MiniBlockCompressionContext { allow_generic_offsets, } } + + /// Returns the padded per-chunk header bytes for the supplied value buffers. + pub(crate) fn chunk_header_bytes(self, value_buffers: u64) -> u64 { + let value_buffer_size_bytes = if self.support_large_chunk { 4 } else { 2 }; + 2_u64 + .saturating_add(self.common_chunk_buffers.saturating_mul(2)) + .saturating_add(value_buffers.saturating_mul(value_buffer_size_bytes)) + .next_multiple_of(8) + } + + pub(crate) fn allows_generic_offsets(self) -> bool { + self.allow_generic_offsets + } } /// Describes the size of a mini-block chunk of data @@ -156,6 +169,21 @@ mod tests { assert_eq!(parse_max_miniblock_values(), 4096); } + #[test] + fn compression_context_matches_padded_chunk_headers() { + let expected = [(0, 8, 16), (1, 8, 16), (2, 16, 16)]; + for (common_buffers, one_value_buffer, two_value_buffers) in expected { + let context = MiniBlockCompressionContext::new(common_buffers, true, true); + assert_eq!(context.chunk_header_bytes(1), one_value_buffer); + assert_eq!(context.chunk_header_bytes(2), two_value_buffers); + } + + let legacy_context = MiniBlockCompressionContext::new(2, false, false); + assert_eq!(legacy_context.chunk_header_bytes(1), 8); + assert_eq!(legacy_context.chunk_header_bytes(2), 16); + assert!(!legacy_context.allows_generic_offsets()); + } + #[test] #[serial] fn test_parse_custom_value() { diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index 39b606a35eb..60c33a52035 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -11,7 +11,7 @@ use arrow_array::OffsetSizeTrait; use byteorder::{ByteOrder, LittleEndian}; -use core::panic; +use prost::Message; use crate::compression::{ BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, @@ -19,7 +19,10 @@ use crate::compression::{ }; use crate::buffer::LanceBuffer; +#[cfg(feature = "bitpacking")] +use crate::compression::block::BITPACK_CHUNK_VALUES; use crate::compression::block::{create_block_decompressor, infer_block_value_type}; +use crate::compression_config::CompressionFieldParams; use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock, VariableWidthBlock}; use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; use crate::encodings::logical::primitive::miniblock::{ @@ -30,21 +33,23 @@ use crate::format::pb21::CompressiveEncoding; use crate::format::pb21::compressive_encoding::Compression; use crate::format::{ProtobufUtils21, pb21}; -use lance_core::utils::bit::pad_bytes_to; use lance_core::{Error, Result}; -#[cfg(test)] mod offsets; +use offsets::{BlockCost, OffsetFamilyCompressor, select_delta_flat_offsets, select_offset_family}; + #[derive(Debug)] pub struct BinaryMiniBlockEncoder { minichunk_size: i64, + generic_offsets: Option, } impl Default for BinaryMiniBlockEncoder { fn default() -> Self { Self { minichunk_size: *AIM_MINICHUNK_SIZE, + generic_offsets: None, } } } @@ -58,101 +63,130 @@ pub static AIM_MINICHUNK_SIZE: std::sync::LazyLock = std::sync::LazyLock::n .unwrap_or(DEFAULT_AIM_MINICHUNK_SIZE) }); -// Make it to support both u32 and u64 -fn chunk_offsets( +#[derive(Debug, Clone, Copy)] +struct BinaryChunkRange { + start_offset_index: usize, + end_offset_index: usize, +} + +fn binary_chunk_ranges( offsets: &[N], - data: &[u8], - alignment: usize, minichunk_size: i64, -) -> (Vec, Vec) { - #[derive(Debug)] - struct ChunkInfo { - chunk_start_offset_in_orig_idx: usize, - chunk_last_offset_in_orig_idx: usize, - // the bytes in every chunk starts at `chunk.bytes_start_offset` - bytes_start_offset: usize, - // every chunk is padded to 8 bytes. - // we need to interpret every chunk as &[u32] so we need it to padded at least to 4 bytes, - // this field can actually be eliminated and I can use `num_bytes` in `MiniBlockChunk` to compute - // the `output_total_bytes`. - padded_chunk_size: usize, +) -> Result> { + if offsets.is_empty() { + return Err(Error::invalid_input( + "Variable-width mini-block offsets cannot be empty", + )); } - - let byte_width: usize = N::get_byte_width(); - let mut chunks_info = vec![]; - let mut chunks = vec![]; + let mut ranges = Vec::new(); let mut last_offset_in_orig_idx = 0; loop { let this_last_offset_in_orig_idx = - search_next_offset_idx(offsets, last_offset_in_orig_idx, minichunk_size); - - let num_values_in_this_chunk = this_last_offset_in_orig_idx - last_offset_in_orig_idx; - let chunk_bytes = offsets[this_last_offset_in_orig_idx] - offsets[last_offset_in_orig_idx]; - let this_chunk_size = - (num_values_in_this_chunk + 1) * byte_width + chunk_bytes.to_usize().unwrap(); - - let padded_chunk_size = this_chunk_size.next_multiple_of(alignment); - debug_assert!(padded_chunk_size > 0); - - let this_chunk_bytes_start_offset = (num_values_in_this_chunk + 1) * byte_width; - chunks_info.push(ChunkInfo { - chunk_start_offset_in_orig_idx: last_offset_in_orig_idx, - chunk_last_offset_in_orig_idx: this_last_offset_in_orig_idx, - bytes_start_offset: this_chunk_bytes_start_offset, - padded_chunk_size, - }); - chunks.push(MiniBlockChunk { - log_num_values: if this_last_offset_in_orig_idx == offsets.len() - 1 { - 0 - } else { - num_values_in_this_chunk.trailing_zeros() as u8 - }, - buffer_sizes: vec![padded_chunk_size as u32], + search_next_offset_idx(offsets, last_offset_in_orig_idx, minichunk_size)?; + ranges.push(BinaryChunkRange { + start_offset_index: last_offset_in_orig_idx, + end_offset_index: this_last_offset_in_orig_idx, }); if this_last_offset_in_orig_idx == offsets.len() - 1 { break; } last_offset_in_orig_idx = this_last_offset_in_orig_idx; } + Ok(ranges) +} - let output_total_bytes = chunks_info - .iter() - .map(|chunk_info| chunk_info.padded_chunk_size) - .sum::(); +// Make it to support both i32 and i64 Arrow offsets. +fn chunk_offsets( + offsets: &[N], + data: &[u8], + alignment: usize, + minichunk_size: i64, +) -> Result<(Vec, Vec)> { + let ranges = binary_chunk_ranges(offsets, minichunk_size)?; + chunk_offsets_with_ranges(offsets, data, alignment, &ranges) +} + +fn chunk_offsets_with_ranges( + offsets: &[N], + data: &[u8], + alignment: usize, + ranges: &[BinaryChunkRange], +) -> Result<(Vec, Vec)> { + let byte_width: usize = N::get_byte_width(); + let mut chunk_sizes = Vec::with_capacity(ranges.len()); + let mut chunks = Vec::with_capacity(ranges.len()); + + for range in ranges { + let num_values = range.end_offset_index - range.start_offset_index; + let chunk_bytes = offsets[range.end_offset_index] - offsets[range.start_offset_index]; + let chunk_bytes = chunk_bytes.to_usize().ok_or_else(|| { + Error::invalid_input("Variable-width mini-block byte length does not fit usize") + })?; + let chunk_size = (num_values + 1) + .checked_mul(byte_width) + .and_then(|offset_bytes| offset_bytes.checked_add(chunk_bytes)) + .ok_or_else(|| { + Error::invalid_input("Variable-width mini-block chunk size overflows usize") + })?; + let padded_chunk_size = chunk_size.next_multiple_of(alignment); + let padded_chunk_size_u32 = u32::try_from(padded_chunk_size).map_err(|_| { + Error::invalid_input(format!( + "Variable-width mini-block chunk has {padded_chunk_size} bytes, exceeding u32::MAX" + )) + })?; + chunk_sizes.push(padded_chunk_size); + chunks.push(MiniBlockChunk { + log_num_values: if range.end_offset_index == offsets.len() - 1 { + 0 + } else { + num_values.trailing_zeros() as u8 + }, + buffer_sizes: vec![padded_chunk_size_u32], + }); + } + let output_total_bytes = chunk_sizes.iter().copied().sum::(); let mut output: Vec = Vec::with_capacity(output_total_bytes); - for chunk in chunks_info { - let this_chunk_offsets: Vec = offsets - [chunk.chunk_start_offset_in_orig_idx..=chunk.chunk_last_offset_in_orig_idx] + for (range, padded_chunk_size) in ranges.iter().zip(chunk_sizes) { + let chunk_output_start = output.len(); + let bytes_start_offset = + (range.end_offset_index - range.start_offset_index + 1) * byte_width; + let bytes_start_offset = N::from_usize(bytes_start_offset).ok_or_else(|| { + Error::invalid_input("Variable-width mini-block offset header does not fit offset type") + })?; + let this_chunk_offsets: Vec = offsets[range.start_offset_index..=range.end_offset_index] .iter() - .map(|offset| { - *offset - offsets[chunk.chunk_start_offset_in_orig_idx] - + N::from_usize(chunk.bytes_start_offset).unwrap() - }) + .map(|offset| *offset - offsets[range.start_offset_index] + bytes_start_offset) .collect(); let this_chunk_offsets = LanceBuffer::reinterpret_vec(this_chunk_offsets); output.extend_from_slice(&this_chunk_offsets); - let start_in_orig = offsets[chunk.chunk_start_offset_in_orig_idx] + let start_in_orig = offsets[range.start_offset_index] .to_usize() - .unwrap(); - let end_in_orig = offsets[chunk.chunk_last_offset_in_orig_idx] - .to_usize() - .unwrap(); + .ok_or_else(|| { + Error::invalid_input("Variable-width mini-block start offset does not fit usize") + })?; + let end_in_orig = offsets[range.end_offset_index].to_usize().ok_or_else(|| { + Error::invalid_input("Variable-width mini-block end offset does not fit usize") + })?; + if start_in_orig > end_in_orig || end_in_orig > data.len() { + return Err(Error::invalid_input(format!( + "Variable-width mini-block byte range {start_in_orig}..{end_in_orig} is invalid for {} bytes", + data.len() + ))); + } output.extend_from_slice(&data[start_in_orig..end_in_orig]); - // pad this chunk to make it align to desired bytes. const PAD_BYTE: u8 = 72; - let pad_len = pad_bytes_to(output.len(), alignment); - - // Compare with usize literal to avoid type mismatch with N - if pad_len > 0_usize { + let encoded_chunk_size = output.len() - chunk_output_start; + let pad_len = padded_chunk_size - encoded_chunk_size; + if pad_len > 0 { output.extend(std::iter::repeat_n(PAD_BYTE, pad_len)); } } - (vec![LanceBuffer::reinterpret_vec(output)], chunks) + Ok((vec![LanceBuffer::reinterpret_vec(output)], chunks)) } // search for the next offset index to cut the values into a chunk. @@ -163,35 +197,35 @@ fn search_next_offset_idx( offsets: &[N], last_offset_idx: usize, minichunk_size: i64, -) -> usize { +) -> Result { // MiniBlockChunk uses `log_num_values == 0` as a sentinel for the final chunk. This means we // must avoid creating 1-value chunks except for the final chunk, even if the configured // `minichunk_size` is too small to fit more than one value. let remaining_values = offsets.len().saturating_sub(last_offset_idx + 1); if remaining_values <= 1 { - return offsets.len() - 1; + return Ok(offsets.len() - 1); } let mut num_values = 2; let mut new_num_values = num_values * 2; loop { if last_offset_idx + new_num_values >= offsets.len() { - let existing_bytes = offsets[offsets.len() - 1] - offsets[last_offset_idx]; - // existing bytes plus the new offset size - let new_size = existing_bytes - + N::from_usize((offsets.len() - last_offset_idx) * N::get_byte_width()).unwrap(); - if new_size.to_i64().unwrap() <= minichunk_size { + let new_size = + checked_variable_chunk_size(offsets, last_offset_idx, offsets.len() - 1)?; + if new_size <= i128::from(minichunk_size) { // case 1: can fit the rest of all data into a miniblock - return offsets.len() - 1; + return Ok(offsets.len() - 1); } else { // case 2: can only fit the last tried `num_values` into a miniblock - return last_offset_idx + num_values; + return Ok(last_offset_idx + num_values); } } - let existing_bytes = offsets[last_offset_idx + new_num_values] - offsets[last_offset_idx]; - let new_size = - existing_bytes + N::from_usize((new_num_values + 1) * N::get_byte_width()).unwrap(); - if new_size.to_i64().unwrap() <= minichunk_size { + let new_size = checked_variable_chunk_size( + offsets, + last_offset_idx, + last_offset_idx + new_num_values, + )?; + if new_size <= i128::from(minichunk_size) { if new_num_values * 2 > *MAX_MINIBLOCK_VALUES as usize { // hit the max number of values limit break; @@ -202,62 +236,593 @@ fn search_next_offset_idx( break; } } - last_offset_idx + num_values + Ok(last_offset_idx + num_values) +} + +fn checked_variable_chunk_size( + offsets: &[N], + start: usize, + end: usize, +) -> Result { + let start_offset = offsets[start] + .to_i64() + .ok_or_else(|| Error::invalid_input("Variable-width chunk start does not fit i64"))?; + let end_offset = offsets[end] + .to_i64() + .ok_or_else(|| Error::invalid_input("Variable-width chunk end does not fit i64"))?; + let value_bytes = i128::from(end_offset) + .checked_sub(i128::from(start_offset)) + .filter(|value_bytes| *value_bytes >= 0) + .ok_or_else(|| { + Error::invalid_input(format!( + "Variable-width offsets decrease between indices {start} and {end}" + )) + })?; + let offset_bytes = (end - start + 1) + .checked_mul(N::get_byte_width()) + .ok_or_else(|| { + Error::invalid_input("Variable-width mini-block offset bytes overflow usize") + })?; + value_bytes + .checked_add(offset_bytes as i128) + .ok_or_else(|| Error::invalid_input("Variable-width mini-block size overflows i128")) +} + +fn validate_variable_offsets( + offsets: &[N], + num_values: u64, + data_len: usize, +) -> Result<()> { + validate_variable_offset_endpoints(offsets, num_values, data_len)?; + let mut previous = None; + for (index, offset) in offsets.iter().enumerate() { + let offset = offset.to_usize().ok_or_else(|| { + Error::invalid_input(format!( + "Variable-width offset at index {index} is negative or does not fit usize" + )) + })?; + if previous.is_some_and(|previous| offset < previous) { + return Err(Error::invalid_input(format!( + "Variable-width offsets decrease at index {index}" + ))); + } + previous = Some(offset); + } + Ok(()) +} + +fn validate_variable_offset_endpoints( + offsets: &[N], + num_values: u64, + data_len: usize, +) -> Result<()> { + let expected_offsets = usize::try_from(num_values) + .ok() + .and_then(|num_values| num_values.checked_add(1)) + .ok_or_else(|| Error::invalid_input("Variable-width offset count overflows usize"))?; + if offsets.len() != expected_offsets { + return Err(Error::invalid_input(format!( + "Variable-width block has {} offsets, expected {expected_offsets}", + offsets.len() + ))); + } + let first = offsets[0].to_usize().ok_or_else(|| { + Error::invalid_input("First variable-width offset is negative or does not fit usize") + })?; + if first != 0 { + return Err(Error::invalid_input(format!( + "Variable-width offsets must start at zero, got {first}" + ))); + } + let last = offsets[offsets.len() - 1].to_usize().ok_or_else(|| { + Error::invalid_input("Final variable-width offset is negative or does not fit usize") + })?; + if last != data_len { + return Err(Error::invalid_input(format!( + "Final variable-width offset {last} does not equal {data_len} data bytes" + ))); + } + Ok(()) +} + +fn legacy_variable_cost( + encoding: &CompressiveEncoding, + offsets: &[N], + ranges: &[BinaryChunkRange], + bits_per_offset: u64, + context: MiniBlockCompressionContext, +) -> Result { + let offset_bytes = usize::try_from(bits_per_offset / 8) + .map_err(|_| Error::invalid_input("Offset width does not fit usize"))?; + ranges + .iter() + .try_fold(encoding.encoded_len() as u64, |cost, range| { + let num_offsets = range.end_offset_index - range.start_offset_index + 1; + let value_bytes = chunk_value_range(offsets, *range)?.len(); + let raw_bytes = num_offsets + .checked_mul(offset_bytes) + .and_then(|bytes| bytes.checked_add(value_bytes)) + .ok_or_else(|| { + Error::invalid_input("Legacy variable chunk size overflows usize") + })?; + Ok(cost + .saturating_add(context.chunk_header_bytes(1)) + .saturating_add((raw_bytes as u64).next_multiple_of(8))) + }) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum GenericOffsetPreflight { + Legacy, + DeltaFlat, + FullSelection, +} + +fn preflight_generic_offsets( + offsets: &[N], + ranges: &[BinaryChunkRange], + bits_per_offset: u64, + legacy_cost: u64, + context: MiniBlockCompressionContext, +) -> Result { + const PROBE_DELTAS_PER_CHUNK: usize = 8; + + let value_bytes = bits_per_offset / 8; + let mut direct_member_lengths = Vec::with_capacity(ranges.len()); + let mut delta_member_lengths = Vec::with_capacity(ranges.len()); + let mut direct_total_values = 0_u64; + let mut delta_total_values = 0_u64; + let mut direct_max = 0_u64; + let mut observed_delta_max = 0_u64; + let mut common_delta = None; + let mut common_delta_range_start = None; + let mut common_delta_range_step = None; + let mut constant_delta_possible = true; + let mut delta_range_possible = true; + + for range in ranges { + let num_deltas = range.end_offset_index - range.start_offset_index; + let num_offsets = num_deltas + 1; + let num_offsets = num_offsets as u64; + let num_deltas = num_deltas as u64; + direct_member_lengths.push(num_offsets); + delta_member_lengths.push(num_deltas); + direct_total_values = direct_total_values + .checked_add(num_offsets) + .ok_or_else(|| Error::invalid_input("Offset family cardinality overflows u64"))?; + delta_total_values = delta_total_values + .checked_add(num_deltas) + .ok_or_else(|| Error::invalid_input("Delta family cardinality overflows u64"))?; + + let start = offsets[range.start_offset_index] + .to_i64() + .and_then(|value| u64::try_from(value).ok()) + .ok_or_else(|| Error::invalid_input("Variable offset does not fit u64"))?; + let end = offsets[range.end_offset_index] + .to_i64() + .and_then(|value| u64::try_from(value).ok()) + .ok_or_else(|| Error::invalid_input("Variable offset does not fit u64"))?; + direct_max = direct_max.max(end.checked_sub(start).ok_or_else(|| { + Error::invalid_input("Variable-width offsets decrease across a chunk") + })?); + + let probe_end = range + .start_offset_index + .saturating_add(PROBE_DELTAS_PER_CHUNK) + .min(range.end_offset_index); + if probe_end - range.start_offset_index < 2 { + delta_range_possible = false; + } + let mut previous_delta = None; + for offset_index in range.start_offset_index..probe_end { + let start = offsets[offset_index] + .to_i64() + .and_then(|value| u64::try_from(value).ok()) + .ok_or_else(|| Error::invalid_input("Variable offset does not fit u64"))?; + let end = offsets[offset_index + 1] + .to_i64() + .and_then(|value| u64::try_from(value).ok()) + .ok_or_else(|| Error::invalid_input("Variable offset does not fit u64"))?; + let delta = end.checked_sub(start).ok_or_else(|| { + Error::invalid_input(format!( + "Variable-width offsets decrease at index {}", + offset_index + 1 + )) + })?; + observed_delta_max = observed_delta_max.max(delta); + match common_delta { + Some(common) => constant_delta_possible &= delta == common, + None => common_delta = Some(delta), + } + if previous_delta.is_none() { + match common_delta_range_start { + Some(common) => delta_range_possible &= delta == common, + None => common_delta_range_start = Some(delta), + } + } + if let Some(previous) = previous_delta { + let step = delta.checked_sub(previous); + match (common_delta_range_step, step) { + (Some(common), Some(step)) => delta_range_possible &= step == common, + (None, Some(step)) => common_delta_range_step = Some(step), + (_, None) => delta_range_possible = false, + } + } + previous_delta = Some(delta); + } + } + + if constant_delta_possible || delta_range_possible { + return Ok(GenericOffsetPreflight::FullSelection); + } + let required_bits = |max: u64| { + if max == 0 { + 1 + } else { + u64::from(u64::BITS - max.leading_zeros()) + } + }; + if !family_bitpacking_cannot_reduce( + direct_total_values, + &direct_member_lengths, + required_bits(direct_max), + bits_per_offset, + ) || !family_bitpacking_cannot_reduce( + delta_total_values, + &delta_member_lengths, + required_bits(observed_delta_max), + bits_per_offset, + ) { + return Ok(GenericOffsetPreflight::FullSelection); + } + + let payload_bytes = delta_member_lengths + .iter() + .map(|num_values| { + num_values + .checked_mul(value_bytes) + .ok_or_else(|| Error::invalid_input("Delta payload size overflows u64")) + }) + .collect::>>()?; + let encoding = ProtobufUtils21::variable( + ProtobufUtils21::delta( + bits_per_offset, + 0, + ProtobufUtils21::flat(bits_per_offset, None), + ), + None, + ); + let generic_cost = + generic_variable_cost(&encoding, true, &payload_bytes, offsets, ranges, context)?; + Ok(if generic_cost < legacy_cost { + GenericOffsetPreflight::DeltaFlat + } else { + GenericOffsetPreflight::Legacy + }) +} + +#[cfg(feature = "bitpacking")] +fn family_bitpacking_cannot_reduce( + total_values: u64, + member_lengths: &[u64], + required_bits: u64, + bits_per_value: u64, +) -> bool { + if required_bits >= bits_per_value { + return true; + } + if total_values <= BITPACK_CHUNK_VALUES { + return false; + } + member_lengths.iter().all(|num_values| { + if *num_values >= BITPACK_CHUNK_VALUES { + return false; + } + let padding_cost = required_bits * (BITPACK_CHUNK_VALUES - num_values); + let tail_savings = (bits_per_value - required_bits) * num_values; + padding_cost >= tail_savings + }) +} + +#[cfg(not(feature = "bitpacking"))] +fn family_bitpacking_cannot_reduce( + _total_values: u64, + _member_lengths: &[u64], + _required_bits: u64, + _bits_per_value: u64, +) -> bool { + true +} + +fn generic_variable_cost( + encoding: &CompressiveEncoding, + has_payload: bool, + payload_bytes: &[u64], + offsets: &[N], + ranges: &[BinaryChunkRange], + context: MiniBlockCompressionContext, +) -> Result { + if payload_bytes.len() != ranges.len() { + return Err(Error::internal(format!( + "Offset family produced {} estimates for {} chunks", + payload_bytes.len(), + ranges.len() + ))); + } + ranges.iter().zip(payload_bytes).try_fold( + encoding.encoded_len() as u64, + |cost, (range, payload)| { + let value_bytes = chunk_value_range(offsets, *range)?.len() as u64; + let payload_cost = if has_payload { + payload.next_multiple_of(8) + } else { + 0 + }; + let value_buffers = 1 + u64::from(has_payload); + Ok(cost + .saturating_add(context.chunk_header_bytes(value_buffers)) + .saturating_add(payload_cost) + .saturating_add(value_bytes.next_multiple_of(8))) + }, + ) +} + +fn chunk_value_range( + offsets: &[N], + range: BinaryChunkRange, +) -> Result> { + let start = offsets[range.start_offset_index] + .to_usize() + .ok_or_else(|| Error::invalid_input("Variable chunk start offset does not fit usize"))?; + let end = offsets[range.end_offset_index] + .to_usize() + .ok_or_else(|| Error::invalid_input("Variable chunk end offset does not fit usize"))?; + if start > end { + return Err(Error::invalid_input( + "Variable chunk offsets are decreasing", + )); + } + Ok(start..end) +} + +fn build_generic_chunks( + data: VariableWidthBlock, + offsets: &[N], + ranges: &[BinaryChunkRange], + family: OffsetFamilyCompressor, + encoding: CompressiveEncoding, +) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + let has_payload = family.has_payload(); + let offset_capacity = + family + .estimated_payload_bytes() + .iter() + .try_fold(0_usize, |total, payload_bytes| { + let payload_bytes = usize::try_from(*payload_bytes).map_err(|_| { + Error::invalid_input("Generic offset payload size does not fit usize") + })?; + total.checked_add(payload_bytes).ok_or_else(|| { + Error::invalid_input("Generic offset payload capacity overflows usize") + }) + })?; + let value_capacity = ranges.iter().try_fold(0_usize, |total, range| { + total + .checked_add(chunk_value_range(offsets, *range)?.len()) + .ok_or_else(|| Error::invalid_input("Variable value capacity overflows usize")) + })?; + let payloads = family.compress_members()?; + let mut offset_data = Vec::with_capacity(offset_capacity); + let mut value_data = Vec::with_capacity(value_capacity); + let mut chunks = Vec::with_capacity(ranges.len()); + + for ((range, payload), is_last) in ranges + .iter() + .zip(payloads) + .zip((0..ranges.len()).map(|index| index + 1 == ranges.len())) + { + let value_range = chunk_value_range(offsets, *range)?; + if value_range.end > data.data.len() { + return Err(Error::invalid_input(format!( + "Variable chunk ends at {}, beyond {} data bytes", + value_range.end, + data.data.len() + ))); + } + let value_bytes = &data.data[value_range]; + let mut buffer_sizes = Vec::with_capacity(1 + usize::from(has_payload)); + if let Some(payload) = payload { + buffer_sizes.push(u32::try_from(payload.len()).map_err(|_| { + Error::invalid_input("Generic offset payload exceeds u32::MAX bytes") + })?); + offset_data.extend_from_slice(&payload); + } + buffer_sizes.push(u32::try_from(value_bytes.len()).map_err(|_| { + Error::invalid_input("Variable chunk value payload exceeds u32::MAX bytes") + })?); + value_data.extend_from_slice(value_bytes); + let num_values = range.end_offset_index - range.start_offset_index; + chunks.push(MiniBlockChunk { + buffer_sizes, + log_num_values: if is_last { + 0 + } else { + num_values.trailing_zeros() as u8 + }, + }); + } + + let mut buffers = Vec::with_capacity(1 + usize::from(has_payload)); + if has_payload { + buffers.push(LanceBuffer::from(offset_data)); + } + buffers.push(LanceBuffer::from(value_data)); + Ok(( + MiniBlockCompressed { + data: buffers, + chunks, + num_values: data.num_values, + }, + encoding, + )) } impl BinaryMiniBlockEncoder { pub fn new(minichunk_size: Option) -> Self { Self { minichunk_size: minichunk_size.unwrap_or(*AIM_MINICHUNK_SIZE), + generic_offsets: None, + } + } + + pub(crate) fn with_generic_offsets( + minichunk_size: Option, + field_params: CompressionFieldParams, + ) -> Self { + Self { + minichunk_size: minichunk_size.unwrap_or(*AIM_MINICHUNK_SIZE), + generic_offsets: Some(field_params), } } // put binary data into chunks, every chunk is less than or equal to `minichunk_size`. // In each chunk, offsets are put first then followed by binary bytes data, each chunk is padded to 8 bytes. // the offsets in the chunk points to the bytes offset in this chunk. - fn chunk_data(&self, data: VariableWidthBlock) -> (MiniBlockCompressed, CompressiveEncoding) { - // TODO: Support compression of offsets - // TODO: Support general compression of data + fn chunk_data( + &self, + data: VariableWidthBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data.bits_per_offset { 32 => { - let offsets = data.offsets.borrow_to_typed_slice::(); - let (buffers, chunks) = - chunk_offsets(offsets.as_ref(), &data.data, 4, self.minichunk_size); - ( - MiniBlockCompressed { - data: buffers, - chunks, - num_values: data.num_values, - }, - ProtobufUtils21::variable(ProtobufUtils21::flat(32, None), None), - ) + let offsets_buffer = data.offsets.clone(); + let offsets = offsets_buffer.borrow_to_typed_slice::(); + self.chunk_typed_data(offsets.as_ref(), data, 32, 4, context) } 64 => { - let offsets = data.offsets.borrow_to_typed_slice::(); - let (buffers, chunks) = - chunk_offsets(offsets.as_ref(), &data.data, 8, self.minichunk_size); - ( - MiniBlockCompressed { - data: buffers, - chunks, - num_values: data.num_values, - }, - ProtobufUtils21::variable(ProtobufUtils21::flat(64, None), None), - ) + let offsets_buffer = data.offsets.clone(); + let offsets = offsets_buffer.borrow_to_typed_slice::(); + self.chunk_typed_data(offsets.as_ref(), data, 64, 8, context) } - _ => panic!("Unsupported bits_per_offset={}", data.bits_per_offset), + _ => Err(Error::invalid_input(format!( + "Unsupported bits_per_offset={}", + data.bits_per_offset + ))), } } + + fn chunk_typed_data( + &self, + offsets: &[N], + data: VariableWidthBlock, + bits_per_offset: u64, + legacy_alignment: usize, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + if context.allows_generic_offsets() + && let Some(field_params) = self.generic_offsets.as_ref() + { + validate_variable_offset_endpoints(offsets, data.num_values, data.data.len())?; + let ranges = binary_chunk_ranges(offsets, self.minichunk_size)?; + let legacy_encoding = + ProtobufUtils21::variable(ProtobufUtils21::flat(bits_per_offset, None), None); + let legacy_cost = + legacy_variable_cost(&legacy_encoding, offsets, &ranges, bits_per_offset, context)?; + let preflight = + preflight_generic_offsets(offsets, &ranges, bits_per_offset, legacy_cost, context)?; + let member_ranges = || { + ranges + .iter() + .map(|range| { + let end = range.end_offset_index.checked_add(1).ok_or_else(|| { + Error::invalid_input("Offset block end overflows usize") + })?; + Ok(range.start_offset_index..end) + }) + .collect::>>() + }; + let offset_block = || FixedWidthDataBlock { + data: data.offsets.clone(), + bits_per_value: bits_per_offset, + num_values: offsets.len() as u64, + block_info: BlockInfo::default(), + }; + let payload_header_bytes = context + .chunk_header_bytes(2) + .saturating_sub(context.chunk_header_bytes(1)); + let block_cost = BlockCost::new(payload_header_bytes, 8); + let family = match preflight { + GenericOffsetPreflight::Legacy => { + validate_variable_offsets(offsets, data.num_values, data.data.len())?; + None + } + GenericOffsetPreflight::DeltaFlat => { + Some(select_delta_flat_offsets(offset_block(), member_ranges()?)?) + } + GenericOffsetPreflight::FullSelection => { + let mut offset_params = field_params.clone(); + // The surrounding mini-block compressor owns general compression. + // Offset selection remains structural and compares exact payload sizes. + offset_params.compression = Some("none".to_string()); + Some(select_offset_family( + offset_block(), + member_ranges()?, + &offset_params, + block_cost, + )?) + } + }; + if let Some(family) = family { + let generic_encoding = ProtobufUtils21::variable(family.encoding().clone(), None); + let generic_cost = generic_variable_cost( + &generic_encoding, + family.has_payload(), + family.estimated_payload_bytes(), + offsets, + &ranges, + context, + )?; + let is_ambiguous_flat = matches!( + family.encoding().compression.as_ref(), + Some(Compression::Flat(_)) + ); + if !is_ambiguous_flat && generic_cost < legacy_cost { + return build_generic_chunks(data, offsets, &ranges, family, generic_encoding); + } + } + let (buffers, chunks) = + chunk_offsets_with_ranges(offsets, &data.data, legacy_alignment, &ranges)?; + return Ok(( + MiniBlockCompressed { + data: buffers, + chunks, + num_values: data.num_values, + }, + legacy_encoding, + )); + } + + validate_variable_offsets(offsets, data.num_values, data.data.len())?; + let (buffers, chunks) = + chunk_offsets(offsets, &data.data, legacy_alignment, self.minichunk_size)?; + Ok(( + MiniBlockCompressed { + data: buffers, + chunks, + num_values: data.num_values, + }, + ProtobufUtils21::variable(ProtobufUtils21::flat(bits_per_offset, None), None), + )) + } } impl MiniBlockCompressor for BinaryMiniBlockEncoder { fn compress( &self, data: DataBlock, - _context: MiniBlockCompressionContext, + context: MiniBlockCompressionContext, ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { - DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)), + DataBlock::VariableWidth(variable_width) => self.chunk_data(variable_width, context), _ => Err(Error::invalid_input_source( format!( "Cannot compress a data block of type {} with BinaryMiniBlockEncoder", @@ -847,7 +1412,7 @@ impl BlockDecompressor for BinaryBlockDecompressor { #[cfg(test)] mod tests { - use super::BinaryMiniBlockDecompressor; + use super::*; use arrow_array::{ ArrayRef, StringArray, builder::{LargeStringBuilder, StringBuilder}, @@ -855,16 +1420,9 @@ mod tests { use arrow_schema::{DataType, Field}; use crate::{ - buffer::LanceBuffer, - compression::MiniBlockDecompressor, constants::{ - COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, - STRUCTURAL_ENCODING_MINIBLOCK, - }, - data::DataBlock, - format::{ - ProtobufUtils21, - pb21::{CompressiveEncoding, compressive_encoding::Compression}, + COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY, STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, }, testing::check_specific_random, }; @@ -872,6 +1430,7 @@ mod tests { use std::{collections::HashMap, sync::Arc, vec}; use crate::{ + compression_config::CompressionFieldParams, testing::{ FnArrayGeneratorProvider, TestCases, check_basic_random, check_round_trip_encoding_of_data, @@ -879,6 +1438,408 @@ mod tests { version::LanceFileVersion, }; + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + + fn decode_binary_miniblocks( + compressed: MiniBlockCompressed, + encoding: &CompressiveEncoding, + ) -> Result> { + let Compression::Variable(variable) = encoding + .compression + .as_ref() + .ok_or_else(|| Error::invalid_input("missing Variable encoding"))? + else { + return Err(Error::invalid_input("expected Variable encoding")); + }; + let decoder = BinaryMiniBlockDecompressor::from_variable(variable)?; + let mut buffer_offsets = vec![0_usize; compressed.data.len()]; + let mut values_seen = 0_u64; + let mut decoded = Vec::with_capacity(compressed.chunks.len()); + for chunk in compressed.chunks { + let num_values = chunk.num_values(values_seen, compressed.num_values); + values_seen += num_values; + let buffers = chunk + .buffer_sizes + .iter() + .zip(compressed.data.iter().zip(&mut buffer_offsets)) + .map(|(size, (buffer, offset))| { + let size = *size as usize; + let chunk = buffer.slice_with_length(*offset, size); + *offset += size; + chunk + }) + .collect(); + let DataBlock::VariableWidth(block) = decoder.decompress(buffers, num_values)? else { + return Err(Error::internal( + "Binary mini-block decoded a non-variable block".to_string(), + )); + }; + decoded.push(block); + } + Ok(decoded) + } + + fn variable_block_u32(lengths: &[usize]) -> VariableWidthBlock { + let mut offsets = Vec::with_capacity(lengths.len() + 1); + let mut data = Vec::new(); + offsets.push(0_i32); + for (index, length) in lengths.iter().copied().enumerate() { + data.extend(std::iter::repeat_n((index % 251) as u8, length)); + offsets.push(i32::try_from(data.len()).unwrap()); + } + VariableWidthBlock { + data: LanceBuffer::from(data), + offsets: LanceBuffer::reinterpret_vec(offsets), + bits_per_offset: 32, + num_values: lengths.len() as u64, + block_info: BlockInfo::default(), + } + } + + fn assert_decoded_value_lengths(decoded: &[VariableWidthBlock], expected: &[usize]) { + let actual = decoded + .iter() + .flat_map(|block| { + let offsets = block.offsets.borrow_to_typed_slice::(); + offsets + .windows(2) + .map(|pair| (pair[1] - pair[0]) as usize) + .collect::>() + }) + .collect::>(); + assert_eq!(actual, expected); + } + + #[test] + fn generic_offsets_use_range_for_fixed_width_values() { + let lengths = vec![3_usize; 2_048]; + let block = variable_block_u32(&lengths); + let encoder = BinaryMiniBlockEncoder::with_generic_offsets( + Some(256), + CompressionFieldParams::default(), + ); + let (compressed, encoding) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding"); + }; + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(Compression::Range(_)) + )); + assert_eq!(compressed.data.len(), 1); + assert!(compressed.chunks.len() > 1); + assert!( + compressed + .chunks + .iter() + .all(|chunk| chunk.buffer_sizes.len() == 1) + ); + + let decoded = decode_binary_miniblocks(compressed, &encoding).unwrap(); + assert_decoded_value_lengths(&decoded, &lengths); + } + + #[test] + fn generic_offsets_use_delta_for_irregular_values() { + let lengths = (0..4_096) + .map(|index| [1_usize, 7, 2, 5][index % 4]) + .collect::>(); + let block = variable_block_u32(&lengths); + let encoder = BinaryMiniBlockEncoder::with_generic_offsets( + Some(1_024), + CompressionFieldParams::default(), + ); + let (compressed, encoding) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding"); + }; + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(Compression::Delta(_)) + )); + assert_eq!(compressed.data.len(), 2); + assert!( + compressed + .chunks + .iter() + .all(|chunk| chunk.buffer_sizes.len() == 2) + ); + + let decoded = decode_binary_miniblocks(compressed, &encoding).unwrap(); + assert_decoded_value_lengths(&decoded, &lengths); + } + + #[test] + fn generic_offsets_use_delta_range_for_increasing_lengths() { + let lengths = (0..4_096) + .map(|index| 4_usize + index % 64) + .collect::>(); + let block = variable_block_u32(&lengths); + let encoder = BinaryMiniBlockEncoder::with_generic_offsets( + Some(4_096), + CompressionFieldParams::default(), + ); + let (compressed, encoding) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding"); + }; + let Some(Compression::Delta(delta)) = variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()) + else { + panic!("expected Delta offsets"); + }; + assert!(matches!( + delta + .deltas + .as_deref() + .and_then(|deltas| deltas.compression.as_ref()), + Some(Compression::Range(_)) + )); + assert_eq!(compressed.data.len(), 1); + let decoded = decode_binary_miniblocks(compressed, &encoding).unwrap(); + assert_decoded_value_lengths(&decoded, &lengths); + } + + #[test] + fn preflight_keeps_legacy_when_delta_flat_does_not_cover_header() { + let lengths = (0..4_096) + .map(|index| [16_usize, 22, 17, 20][index % 4]) + .collect::>(); + let block = variable_block_u32(&lengths); + let offsets = block.offsets.borrow_to_typed_slice::(); + let ranges = binary_chunk_ranges(offsets.as_ref(), DEFAULT_AIM_MINICHUNK_SIZE).unwrap(); + let encoding = ProtobufUtils21::variable(ProtobufUtils21::flat(32, None), None); + let context = MiniBlockCompressionContext::new(0, true, true); + let legacy_cost = + legacy_variable_cost(&encoding, offsets.as_ref(), &ranges, 32, context).unwrap(); + assert_eq!( + preflight_generic_offsets(offsets.as_ref(), &ranges, 32, legacy_cost, context).unwrap(), + GenericOffsetPreflight::Legacy + ); + } + + #[test] + fn generic_offsets_keep_smaller_legacy_container() { + let block = variable_block_u32(&[1, 7, 2, 5]); + let encoder = BinaryMiniBlockEncoder::with_generic_offsets( + Some(4_096), + CompressionFieldParams::default(), + ); + let (compressed, encoding) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding"); + }; + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(Compression::Flat(_)) + )); + assert_eq!(compressed.data.len(), 1); + assert_eq!(compressed.chunks[0].buffer_sizes.len(), 1); + } + + #[test] + fn generic_offsets_support_metadata_only_empty_values() { + let lengths = vec![0_usize; 1_024]; + let block = variable_block_u32(&lengths); + let encoder = BinaryMiniBlockEncoder::with_generic_offsets( + Some(256), + CompressionFieldParams::default(), + ); + let (compressed, encoding) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding"); + }; + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(Compression::Constant(_)) + )); + assert_eq!(compressed.data.len(), 1); + let decoded = decode_binary_miniblocks(compressed, &encoding).unwrap(); + assert_decoded_value_lengths(&decoded, &lengths); + } + + #[test] + fn generic_offsets_support_u64_range() { + let num_values = 512_usize; + let offsets = (0..=num_values) + .map(|index| (index * 2) as i64) + .collect::>(); + let block = VariableWidthBlock { + data: LanceBuffer::from(vec![1_u8; num_values * 2]), + offsets: LanceBuffer::reinterpret_vec(offsets), + bits_per_offset: 64, + num_values: num_values as u64, + block_info: BlockInfo::default(), + }; + let encoder = BinaryMiniBlockEncoder::with_generic_offsets( + Some(256), + CompressionFieldParams::default(), + ); + let (compressed, encoding) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding"); + }; + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(Compression::Range(_)) + )); + let decoded = decode_binary_miniblocks(compressed, &encoding).unwrap(); + assert!(decoded.iter().all(|block| block.bits_per_offset == 64)); + } + + #[test] + fn legacy_offsets_remain_interleaved_flat() { + let block = variable_block_u32(&vec![3_usize; 128]); + let encoder = BinaryMiniBlockEncoder::new(Some(256)); + let (compressed, encoding) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding"); + }; + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(Compression::Flat(_)) + )); + assert_eq!(compressed.data.len(), 1); + assert!( + compressed + .chunks + .iter() + .all(|chunk| chunk.buffer_sizes.len() == 1) + ); + } + + #[test] + fn legacy_u32_interleaved_bytes_are_stable() { + let block = variable_block_u32(&[3, 3]); + let encoder = BinaryMiniBlockEncoder::new(Some(4_096)); + let (compressed, _) = encoder + .compress(DataBlock::VariableWidth(block), miniblock_context()) + .unwrap(); + assert_eq!(compressed.chunks.len(), 1); + assert_eq!(compressed.chunks[0].buffer_sizes, [20]); + + let mut expected = Vec::new(); + expected.extend_from_slice(&12_i32.to_le_bytes()); + expected.extend_from_slice(&15_i32.to_le_bytes()); + expected.extend_from_slice(&18_i32.to_le_bytes()); + expected.extend_from_slice(&[0, 0, 0, 1, 1, 1]); + expected.extend_from_slice(&[72, 72]); + assert_eq!(compressed.data[0].as_ref(), expected); + } + + #[test] + fn generic_offsets_reject_wrong_buffer_count_and_bounds() { + let variable = pb21::Variable { + offsets: Some(Box::new(ProtobufUtils21::range(32, 0, 3))), + values: None, + }; + let decoder = BinaryMiniBlockDecompressor::from_variable(&variable).unwrap(); + let error = decoder + .decompress(vec![LanceBuffer::empty(), LanceBuffer::empty()], 2) + .unwrap_err(); + assert!(error.to_string().contains("requires 1 buffers")); + + let error = decoder + .decompress(vec![LanceBuffer::from(vec![0_u8; 5])], 2) + .unwrap_err(); + assert!(error.to_string().contains("final offset 6")); + } + + #[test] + fn generic_offsets_only_skip_full_scan_for_monotonic_codecs() { + let range = ProtobufUtils21::range(32, 0, 3); + let delta = ProtobufUtils21::delta(32, 0, ProtobufUtils21::flat(32, None)); + let rle = ProtobufUtils21::rle( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::constant(None), + ); + assert_eq!( + offset_validation(range.compression.as_ref().unwrap()), + OffsetValidation::Endpoints + ); + assert_eq!( + offset_validation(delta.compression.as_ref().unwrap()), + OffsetValidation::Endpoints + ); + assert_eq!( + offset_validation(rle.compression.as_ref().unwrap()), + OffsetValidation::Full + ); + + let offsets = FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_u32, 5, 3]), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::default(), + }; + let error = + validate_decoded_offsets(&offsets, BlockValueType::UInt32, 3, OffsetValidation::Full) + .unwrap_err(); + assert!(error.to_string().contains("decrease at index 2")); + } + + #[rstest] + #[case::range([16_usize; 4], "range")] + #[case::delta([4_usize, 10, 5, 8], "delta")] + #[test_log::test(tokio::test)] + async fn generic_offsets_support_scan_range_take( + #[case] lengths: [usize; 4], + #[case] expected_encoding: &str, + ) { + let values = StringArray::from_iter_values((0..10_000).map(|index| { + let len = lengths[index % lengths.len()]; + format!("{index:04x}{}", "x".repeat(len - 4)) + })); + let metadata = HashMap::from([ + ( + STRUCTURAL_ENCODING_META_KEY.to_string(), + STRUCTURAL_ENCODING_MINIBLOCK.to_string(), + ), + (COMPRESSION_META_KEY.to_string(), "none".to_string()), + (DICT_DIVISOR_META_KEY.to_string(), "100000".to_string()), + ]); + let test_cases = TestCases::basic() + .with_min_file_version(LanceFileVersion::V2_3) + .with_expected_encoding(expected_encoding); + check_round_trip_encoding_of_data(vec![Arc::new(values)], &test_cases, metadata).await; + } + fn variable_miniblock_decoder(offsets: CompressiveEncoding) -> BinaryMiniBlockDecompressor { let encoding = ProtobufUtils21::variable(offsets, None); let Compression::Variable(variable) = encoding.compression.as_ref().unwrap() else { diff --git a/rust/lance-encoding/src/encodings/physical/constant.rs b/rust/lance-encoding/src/encodings/physical/constant.rs index 67f6aa6fa3d..9084567e9ee 100644 --- a/rust/lance-encoding/src/encodings/physical/constant.rs +++ b/rust/lance-encoding/src/encodings/physical/constant.rs @@ -3,12 +3,11 @@ //! 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, BlockValueType, FixedPerValueDecompressor, require_no_block_payload, + BlockCompressor, BlockDecompressor, BlockValueType, FixedPerValueDecompressor, + block::validate_fixed_payload_len, require_no_block_payload, }, data::{AllNullDataBlock, BlockInfo, ConstantDataBlock, DataBlock, FixedWidthDataBlock}, encodings::physical::try_vec_with_capacity, @@ -62,21 +61,18 @@ impl FixedPerValueDecompressor for ConstantDecompressor { } /// 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 { diff --git a/rust/lance-encoding/src/encodings/physical/delta.rs b/rust/lance-encoding/src/encodings/physical/delta.rs index fba69e4d6c3..5fade5f26b5 100644 --- a/rust/lance-encoding/src/encodings/physical/delta.rs +++ b/rust/lance-encoding/src/encodings/physical/delta.rs @@ -12,7 +12,6 @@ use crate::{ use lance_core::{Error, Result}; /// Converts a non-decreasing u32/u64 block into adjacent differences. -#[cfg(test)] pub(crate) fn encode_deltas( data: FixedWidthDataBlock, expected_base: u64, diff --git a/rust/lance-encoding/src/encodings/physical/dictionary.rs b/rust/lance-encoding/src/encodings/physical/dictionary.rs index 5713936fd69..50225b55e6e 100644 --- a/rust/lance-encoding/src/encodings/physical/dictionary.rs +++ b/rust/lance-encoding/src/encodings/physical/dictionary.rs @@ -5,26 +5,20 @@ #[cfg(test)] use std::cell::Cell; -#[cfg(test)] use std::{collections::BTreeMap, sync::Arc}; use crate::{ buffer::LanceBuffer, compression::{ - BlockDecompressor, BlockValueType, - block::{fixed_block, read_unsigned_values, validate_fixed_payload_len}, + BlockCompressor, BlockDecompressor, BlockValueType, + block::{ + fixed_block, fixed_from_u64_values, read_unsigned_values, validate_fixed_payload_len, + visit_unsigned_values, + }, }, - data::DataBlock, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, encodings::physical::try_vec_with_capacity, }; -#[cfg(test)] -use crate::{ - compression::{ - BlockCompressor, - block::{fixed_from_u64_values, visit_unsigned_values}, - }, - data::{BlockInfo, FixedWidthDataBlock}, -}; use lance_core::{Error, Result}; pub(crate) const BLOCK_FRAME_BYTES: u64 = 16; @@ -35,7 +29,6 @@ thread_local! { } /// Dictionary compressor that owns its indices and items compressors. -#[cfg(test)] #[derive(Debug)] pub(crate) struct DictionaryBlockCompressor { value_type: BlockValueType, @@ -44,7 +37,6 @@ pub(crate) struct DictionaryBlockCompressor { items: Box, } -#[cfg(test)] impl DictionaryBlockCompressor { pub(crate) fn new( value_type: BlockValueType, @@ -61,7 +53,6 @@ impl DictionaryBlockCompressor { } } -#[cfg(test)] impl BlockCompressor for DictionaryBlockCompressor { fn compress(&self, data: DataBlock) -> Result> { #[cfg(test)] @@ -313,7 +304,6 @@ fn append_items( Ok(()) } -#[cfg(test)] fn try_frame(indices_payload_bytes: usize, items_payload_bytes: usize) -> Result> { let capacity = (BLOCK_FRAME_BYTES as usize) .checked_add(indices_payload_bytes) diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index ccf17e439c9..ff66363affd 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -3,13 +3,12 @@ use log::trace; -#[cfg(test)] -use crate::compression::BlockCompressor; use crate::{ Result, buffer::LanceBuffer, compression::{ - BlockDecompressor, BlockValueType, MiniBlockDecompressor, require_block_payload, + BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, + require_block_payload, }, data::DataBlock, encodings::{ @@ -22,7 +21,6 @@ use crate::{ }; use lance_core::Error; -#[cfg(test)] pub(crate) fn compress_block( compression: CompressionConfig, payload: &[u8], @@ -45,21 +43,18 @@ pub(crate) fn decompress_block_exact( } /// General-purpose block compressor that owns its child block compressor. -#[cfg(test)] #[derive(Debug)] pub(crate) struct GeneralBlockCompressor { child: Box, compression: CompressionConfig, } -#[cfg(test)] impl GeneralBlockCompressor { pub(crate) fn new(child: Box, compression: CompressionConfig) -> Self { Self { child, compression } } } -#[cfg(test)] impl BlockCompressor for GeneralBlockCompressor { fn compress(&self, data: DataBlock) -> Result> { let payload = self.child.compress(data)?.ok_or_else(|| { diff --git a/rust/lance-encoding/src/encodings/physical/range.rs b/rust/lance-encoding/src/encodings/physical/range.rs index 23650c7e66c..21a06c111ec 100644 --- a/rust/lance-encoding/src/encodings/physical/range.rs +++ b/rust/lance-encoding/src/encodings/physical/range.rs @@ -4,11 +4,9 @@ //! Metadata-only arithmetic range encoding for unsigned block sequences. use super::try_vec_with_capacity; -#[cfg(test)] -use crate::compression::BlockCompressor; use crate::{ buffer::LanceBuffer, - compression::{BlockDecompressor, require_no_block_payload}, + compression::{BlockCompressor, BlockDecompressor, require_no_block_payload}, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, }; use lance_core::{Error, Result}; @@ -53,7 +51,6 @@ pub(crate) fn checked_range_last( } /// Validates an input block against the selected arithmetic range codec. -#[cfg(test)] #[derive(Debug)] pub(crate) struct RangeEncoder { bits_per_value: u64, @@ -61,7 +58,6 @@ pub(crate) struct RangeEncoder { step: u64, } -#[cfg(test)] impl RangeEncoder { pub(crate) fn new(bits_per_value: u64, start: u64, step: u64) -> Self { Self { @@ -72,7 +68,6 @@ impl RangeEncoder { } } -#[cfg(test)] impl BlockCompressor for RangeEncoder { fn compress(&self, data: DataBlock) -> Result> { let DataBlock::FixedWidth(data) = data else { @@ -200,7 +195,6 @@ pub(crate) fn materialize_validated_range( })) } -#[cfg(test)] fn checked_values( data: &FixedWidthDataBlock, ) -> Result> { diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index b3d1f91757e..0e7a9ecf58e 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -58,11 +58,12 @@ use arrow_buffer::{ArrowNativeType, ScalarBuffer}; use log::trace; use crate::buffer::LanceBuffer; -#[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}, + block::{ + fixed_block, fixed_from_u64_values, read_unsigned_values, validate_fixed_payload_len, + visit_unsigned_values, + }, require_block_payload, }; use crate::data::DataBlock; @@ -1932,7 +1933,6 @@ impl RleDecompressor { 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, @@ -1941,7 +1941,6 @@ pub(crate) struct BlockRleCompressor { run_lengths: Box, } -#[cfg(test)] impl BlockRleCompressor { pub(crate) fn new( value_type: BlockValueType, @@ -1958,7 +1957,6 @@ impl BlockRleCompressor { } } -#[cfg(test)] impl BlockCompressor for BlockRleCompressor { fn compress(&self, data: DataBlock) -> Result> { let DataBlock::FixedWidth(data) = data else { @@ -1987,7 +1985,6 @@ impl BlockCompressor for BlockRleCompressor { } } -#[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) @@ -2323,7 +2320,6 @@ pub(crate) fn expand_block( Ok(fixed_block(value_type, num_values, output)) } -#[cfg(test)] pub(crate) fn materialize_block( data: &FixedWidthDataBlock, value_type: BlockValueType, diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index 3b7f2365393..7789e7cfe3d 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -474,20 +474,17 @@ impl BlockCompressor for ValueEncoder { } /// 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 {