diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 1cb83d581e9..d10137309c4 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -585,6 +585,8 @@ on a per-value basis. We use ☑️ to mark a technique that is applied on a per | Flat | ✅ (2.1) | ✅ (2.1) | ✅ (2.1) | | Variable | ✅ (2.1) | ✅ (2.1) | ✅ (2.1) | | Constant | ✅ (2.1) | ❓ | ❓ | +| Range | ✅ (2.3) | ❌ | ❓ | +| Delta | ✅ (2.3) | ❌ | ❓ | | Bitpacking | ✅ (2.1) | ❓ | ✅ (2.1) | | Fsst | ❓ | ✅ (2.1) | ✅ (2.1) | | Rle | ✅ (2.2) | ❌ | ✅ (2.1) | @@ -594,6 +596,28 @@ on a per-value basis. We use ☑️ to mark a technique that is applied on a per In the following sections we will describe each technique in a bit more detail and explain how it is utilized in various contexts. +### Generic Block Sequences + +Starting in Lance 2.3, block compression can describe unsigned `u32` and `u64` sequences with a bounded, +zero-or-one-payload codec tree. The containing layout supplies the value type and cardinality. + +`Range` is metadata-only. It stores the unsigned width, first value, and positive step. The value at index `i` +is `start + step * i`; readers reject cardinalities below two, overflow, and widths other than 32 or 64 bits. + +```protobuf +%%% proto.message.Range %%% +``` + +`Delta` stores the first value inline. Its child describes the `n - 1` non-negative adjacent differences and +determines whether the Delta tree has a payload. Zero differences are valid. Readers reconstruct the sequence +with checked prefix sums. + +```protobuf +%%% proto.message.Delta %%% +``` + +Lance 2.0 through 2.2 writers do not emit either encoding. + ### Flat Flat compression is the uncompressed representation of fixed-width data. There is a single buffer of data diff --git a/protos/encodings_v2_1.proto b/protos/encodings_v2_1.proto index 51427332063..77f261768de 100644 --- a/protos/encodings_v2_1.proto +++ b/protos/encodings_v2_1.proto @@ -426,6 +426,39 @@ message Constant { optional bytes value = 1; } +// A metadata-only arithmetic sequence of unsigned fixed-width values. +// +// The value at index i is `start + step * i`. The container supplies the +// number of values. Writers use this encoding only when step is positive and +// the final value fits in the declared bit width. +// +// The input is a u32 or u64 fixed-width data block. +// There is no output buffer. +message Range { + // The width of each decompressed value. Must be 32 or 64. + uint64 uncompressed_bits_per_value = 1; + // The first value in the sequence. + uint64 start = 2; + // The positive difference between adjacent values. + uint64 step = 3; +} + +// A non-decreasing unsigned fixed-width sequence represented by its first +// value and the differences between adjacent values. +// +// The child contains one fewer value than the input. The container supplies +// the input cardinality. Delta has a payload exactly when its child has one. +// +// The input is a u32 or u64 fixed-width data block. +message Delta { + // The width of each decompressed value. Must be 32 or 64. + uint64 uncompressed_bits_per_value = 1; + // The first uncompressed value. + uint64 base = 2; + // Compression applied to the adjacent differences. + CompressiveEncoding deltas = 3; +} + // A compression scheme in which a single fixed-width block is "packed" into // a smaller fixed-width block values where each value has fewer bits. // @@ -631,5 +664,7 @@ message CompressiveEncoding { FixedSizeList fixed_size_list = 11; PackedStruct packed_struct = 12; VariablePackedStruct variable_packed_struct = 13; + Range range = 14; + Delta delta = 15; } } diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 6cb173216b0..ef832726ae8 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -1209,6 +1209,11 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Compression::OutOfLineBitpacking(_) => Err(Error::not_supported_source( "this runtime was not built with bitpacking support".into(), )), + Compression::Range(_) | Compression::Delta(_) => { + let value_type = block::infer_block_value_type(description)?; + block::create_block_decompressor(description, value_type) + .map(|(decompressor, _)| decompressor) + } Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)), Compression::Variable(variable) => { let offsets = variable.offsets.as_deref().ok_or_else(|| { @@ -1443,6 +1448,8 @@ fn compression_name(compression: &Compression) -> &'static str { Compression::FixedSizeList(_) => "fixed-size list", Compression::VariablePackedStruct(_) => "variable packed struct", Compression::Rle(_) => "rle", + Compression::Range(_) => "range", + Compression::Delta(_) => "delta", } } diff --git a/rust/lance-encoding/src/compression/block.rs b/rust/lance-encoding/src/compression/block.rs index 4ca9e9785db..792ca41937c 100644 --- a/rust/lance-encoding/src/compression/block.rs +++ b/rust/lance-encoding/src/compression/block.rs @@ -51,7 +51,6 @@ impl BlockValueType { (self.bits_per_value() / 8) as usize } - #[cfg(test)] pub(crate) fn max_value(self) -> u64 { match self { Self::UInt8 => u8::MAX as u64, @@ -67,7 +66,9 @@ pub(crate) mod fixed; #[cfg(test)] pub(crate) use factory::encode_scalar; -pub(crate) use factory::{create_block_decompressor, validate_fixed_payload_len}; +pub(crate) use factory::{ + create_block_decompressor, 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)] diff --git a/rust/lance-encoding/src/compression/block/factory.rs b/rust/lance-encoding/src/compression/block/factory.rs index 55bd03376d7..525049cdce0 100644 --- a/rust/lance-encoding/src/compression/block/factory.rs +++ b/rust/lance-encoding/src/compression/block/factory.rs @@ -13,7 +13,9 @@ use crate::{ encodings::physical::{ block::{CompressionConfig, CompressionScheme}, constant::ConstantBlockDecompressor, + delta::DeltaDecompressor, general::GenericGeneralBlockDecompressor, + range::RangeDecompressor, rle::{BlockRleDecompressor, BlockRunCount, MetadataRunLengths}, value::FixedWidthBlockDecompressor, }, @@ -198,6 +200,65 @@ fn create_inner( false, )) } + Compression::Range(range) => { + validate_declared_bits(range.uncompressed_bits_per_value, expected_type, "Range")?; + if !matches!( + expected_type, + BlockValueType::UInt32 | BlockValueType::UInt64 + ) { + return Err(Error::invalid_input(format!( + "Range only supports 32 or 64-bit values, got {expected_bits}" + ))); + } + if range.start > expected_type.max_value() { + return Err(Error::invalid_input(format!( + "Range start {} exceeds the {expected_bits}-bit value range", + range.start + ))); + } + if range.step == 0 { + return Err(Error::invalid_input("Range step must be positive")); + } + Ok(( + Box::new(RangeDecompressor::new( + expected_bits, + range.start, + range.step, + )), + false, + )) + } + Compression::Delta(delta) => { + if position != Position::Root { + return Err(Error::invalid_input( + "Delta is not supported as a block codec child", + )); + } + validate_declared_bits(delta.uncompressed_bits_per_value, expected_type, "Delta")?; + if !matches!( + expected_type, + BlockValueType::UInt32 | BlockValueType::UInt64 + ) { + return Err(Error::invalid_input(format!( + "Delta only supports 32 or 64-bit values, got {expected_bits}" + ))); + } + if delta.base > expected_type.max_value() { + return Err(Error::invalid_input(format!( + "Delta base {} exceeds the {expected_bits}-bit value range", + delta.base + ))); + } + let child = delta.deltas.as_deref().ok_or_else(|| { + Error::invalid_input("Delta is missing its deltas child encoding") + })?; + let (child, child_has_payload) = + create_inner(child, expected_type, Position::Child, false)?; + Ok(( + Box::new(DeltaDecompressor::new(expected_bits, delta.base, child)), + child_has_payload, + )) + } Compression::InlineBitpacking(bitpacking) => { validate_declared_bits( bitpacking.uncompressed_bits_per_value, @@ -332,6 +393,17 @@ fn create_inner( )?; Some(MetadataRunLengths::Constant(value)) } + Some(Compression::Range(range)) => { + validate_declared_bits( + range.uncompressed_bits_per_value, + run_length_type, + "RLE run lengths Range", + )?; + Some(MetadataRunLengths::Range { + start: range.start, + step: range.step, + }) + } _ => None, }; let run_count = if let Some(metadata) = metadata_run_lengths { @@ -444,6 +516,11 @@ fn validate_compression_config( Ok(CompressionConfig::new(scheme, compression.level)) } +/// Infers the unsigned fixed-width result of a bounded block descriptor. +pub fn infer_block_value_type(encoding: &CompressiveEncoding) -> Result { + infer_inner(encoding, Position::Root) +} + fn infer_inner(encoding: &CompressiveEncoding, position: Position) -> Result { let compression = encoding .compression @@ -465,6 +542,10 @@ fn infer_inner(encoding: &CompressiveEncoding, position: Position) -> Result { BlockValueType::from_bits(bitpacking.uncompressed_bits_per_value) } + Compression::Range(range) => BlockValueType::from_bits(range.uncompressed_bits_per_value), + Compression::Delta(delta) if position == Position::Root => { + BlockValueType::from_bits(delta.uncompressed_bits_per_value) + } Compression::General(general) if position == Position::Root => infer_inner( general .values @@ -500,6 +581,8 @@ fn compression_name(compression: &Compression) -> &'static str { Compression::FixedSizeList(_) => "fixed-size list", Compression::PackedStruct(_) => "packed struct", Compression::VariablePackedStruct(_) => "variable packed struct", + Compression::Range(_) => "range", + Compression::Delta(_) => "delta", } } @@ -530,4 +613,19 @@ mod tests { let error = create_block_decompressor(&flat, BlockValueType::UInt64).unwrap_err(); assert!(error.to_string().contains("expected 64")); } + + #[test] + fn validates_range_cardinality_when_decoding() { + let encoding = crate::format::ProtobufUtils21::range(32, u32::MAX as u64, 1); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt32).unwrap(); + assert!(!has_payload); + let error = decoder.decompress(None, 2).unwrap_err(); + assert!(error.to_string().contains("exceeds u32::MAX")); + } + + #[test] + fn range_helper_rejects_overflow() { + assert!(crate::encodings::physical::range::checked_range_last(64, u64::MAX, 1, 2).is_err()); + } } diff --git a/rust/lance-encoding/src/compression/block/tests.rs b/rust/lance-encoding/src/compression/block/tests.rs index 372cf56a2b1..bd3b566396c 100644 --- a/rust/lance-encoding/src/compression/block/tests.rs +++ b/rust/lance-encoding/src/compression/block/tests.rs @@ -8,7 +8,7 @@ use crate::{ compression::BlockCompressor, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, encodings::physical::{ - constant::ConstantBlockCompressor, rle::BlockRleCompressor, + constant::ConstantBlockCompressor, range::RangeEncoder, rle::BlockRleCompressor, value::FixedWidthBlockCompressor, }, format::{ProtobufUtils21, pb21::CompressiveEncoding}, @@ -30,6 +30,13 @@ fn decoded_u64(block: DataBlock) -> Vec { block.data.borrow_to_typed_slice::().to_vec() } +fn decoded_u32(block: DataBlock) -> Vec { + let DataBlock::FixedWidth(block) = block else { + panic!("expected fixed-width output"); + }; + block.data.borrow_to_typed_slice::().to_vec() +} + #[test] fn scalar_encoding_is_little_endian_and_bounded() { assert_eq!( @@ -73,6 +80,12 @@ fn metadata_compressors_validate_and_decode() { )), ProtobufUtils21::constant(Some(7_u64.to_le_bytes().to_vec().into())), ); + let values = (0..32_u64).map(|value| 5 + value * 3).collect::>(); + round_trip_u64( + &values, + Box::new(RangeEncoder::new(64, 5, 3)), + ProtobufUtils21::range(64, 5, 3), + ); } #[test] @@ -217,6 +230,40 @@ fn constant_cardinality_contract_is_checked_at_decode() { ); } +#[test] +fn range_checks_cardinality_and_overflow() { + let (decoder, has_payload) = + create_block_decompressor(&ProtobufUtils21::range(32, 3, 5), BlockValueType::UInt32) + .unwrap(); + assert!(!has_payload); + assert_eq!( + decoded_u32(decoder.decompress(None, 4).unwrap()), + vec![3, 8, 13, 18] + ); + assert!(decoder.decompress(None, 1).is_err()); + + let (overflowing, has_payload) = create_block_decompressor( + &ProtobufUtils21::range(32, u32::MAX as u64, 1), + BlockValueType::UInt32, + ) + .unwrap(); + assert!(!has_payload); + assert!(overflowing.decompress(None, 2).is_err()); +} + +#[test] +fn delta_checks_prefix_sum_overflow() { + let encoding = ProtobufUtils21::delta( + 64, + u64::MAX, + ProtobufUtils21::constant(Some(1_u64.to_le_bytes().to_vec().into())), + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(!has_payload); + assert!(decoder.decompress(None, 2).is_err()); +} + #[test] fn metadata_only_rle_round_trip() { let encoding = ProtobufUtils21::rle( @@ -233,6 +280,22 @@ fn metadata_only_rle_round_trip() { assert!(decoder.decompress(None, 5).is_err()); } +#[test] +fn metadata_range_rle_round_trip() { + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::range(64, 10, 1), + ProtobufUtils21::range(32, 1, 1), + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(!has_payload); + assert_eq!( + decoded_u64(decoder.decompress(None, 6).unwrap()), + vec![10, 11, 11, 12, 12, 12] + ); + assert!(decoder.decompress(None, 7).is_err()); +} + #[test] fn rle_framing_is_fallible() { let encoding = ProtobufUtils21::rle( diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs index c0f6901d962..d38c32d7ecb 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs @@ -1386,8 +1386,10 @@ impl SparseStructuralScheduler { .filter_map(|field| field.value.as_ref()), ); } + Compression::Delta(delta) => stack.extend(delta.deltas.as_deref()), Compression::Flat(_) | Compression::Constant(_) + | Compression::Range(_) | Compression::InlineBitpacking(_) => {} } } @@ -1535,7 +1537,9 @@ impl SparseStructuralScheduler { Compression::Constant(_) | Compression::OutOfLineBitpacking(_) | Compression::Dictionary(_) - | Compression::VariablePackedStruct(_) => Err(Error::invalid_input_source( + | Compression::VariablePackedStruct(_) + | Compression::Range(_) + | Compression::Delta(_) => Err(Error::invalid_input_source( "Sparse value compression uses an unsupported mini-block encoding".into(), )), } diff --git a/rust/lance-encoding/src/encodings/physical.rs b/rust/lance-encoding/src/encodings/physical.rs index ea13dad7970..849b720e4d0 100644 --- a/rust/lance-encoding/src/encodings/physical.rs +++ b/rust/lance-encoding/src/encodings/physical.rs @@ -9,9 +9,11 @@ pub mod bitpacking; pub mod block; pub mod byte_stream_split; pub mod constant; +pub mod delta; pub mod fsst; pub mod general; pub mod packed; +pub mod range; pub mod rle; pub mod value; diff --git a/rust/lance-encoding/src/encodings/physical/delta.rs b/rust/lance-encoding/src/encodings/physical/delta.rs new file mode 100644 index 00000000000..fba69e4d6c3 --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/delta.rs @@ -0,0 +1,277 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Delta transform for non-decreasing unsigned block sequences. + +use super::try_vec_with_capacity; +use crate::{ + buffer::LanceBuffer, + compression::BlockDecompressor, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, +}; +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, +) -> Result { + if !matches!(data.bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Delta only supports 32 or 64-bit values, got {}", + data.bits_per_value + ))); + } + if data.num_values < 2 { + return Err(Error::invalid_input(format!( + "Delta requires at least 2 values, got {}", + data.num_values + ))); + } + + match data.bits_per_value { + 32 => { + let values = checked_values::(&data, "Delta")?; + if u64::from(values[0]) != expected_base { + return Err(Error::invalid_input(format!( + "Delta base mismatch: codec expects {expected_base}, input starts with {}", + values[0] + ))); + } + let mut deltas = Vec::with_capacity(values.len() - 1); + for (index, pair) in values.windows(2).enumerate() { + deltas.push(pair[1].checked_sub(pair[0]).ok_or_else(|| { + Error::invalid_input(format!( + "Delta input decreases at index {}: {} -> {}", + index + 1, + pair[0], + pair[1] + )) + })?); + } + Ok(FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(deltas), + num_values: data.num_values - 1, + block_info: BlockInfo::default(), + }) + } + 64 => { + let values = checked_values::(&data, "Delta")?; + if values[0] != expected_base { + return Err(Error::invalid_input(format!( + "Delta base mismatch: codec expects {expected_base}, input starts with {}", + values[0] + ))); + } + let mut deltas = Vec::with_capacity(values.len() - 1); + for (index, pair) in values.windows(2).enumerate() { + deltas.push(pair[1].checked_sub(pair[0]).ok_or_else(|| { + Error::invalid_input(format!( + "Delta input decreases at index {}: {} -> {}", + index + 1, + pair[0], + pair[1] + )) + })?); + } + Ok(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(deltas), + num_values: data.num_values - 1, + block_info: BlockInfo::default(), + }) + } + _ => unreachable!("delta width was validated above"), + } +} + +/// Reconstructs a delta sequence after its child has been decoded. +#[derive(Debug)] +pub(crate) struct DeltaDecompressor { + bits_per_value: u64, + base: u64, + child: Box, +} + +impl DeltaDecompressor { + pub(crate) fn new(bits_per_value: u64, base: u64, child: Box) -> Self { + Self { + bits_per_value, + base, + child, + } + } +} + +impl BlockDecompressor for DeltaDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + if num_values < 2 { + return Err(Error::invalid_input(format!( + "Delta requires at least 2 values, got {num_values}" + ))); + } + if self.bits_per_value == 32 && self.base > u32::MAX as u64 { + return Err(Error::invalid_input(format!( + "Delta base {} exceeds u32::MAX", + self.base + ))); + } + if !matches!(self.bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Delta only supports 32 or 64-bit values, got {}", + self.bits_per_value + ))); + } + + let child = self.child.decompress(data, num_values - 1)?; + reconstruct_deltas(child, self.bits_per_value, self.base, num_values) + } +} + +pub(crate) fn reconstruct_deltas( + child: DataBlock, + bits_per_value: u64, + base: u64, + num_values: u64, +) -> Result { + if num_values < 2 { + return Err(Error::invalid_input(format!( + "Delta requires at least 2 values, got {num_values}" + ))); + } + let DataBlock::FixedWidth(child) = child else { + return Err(Error::invalid_input( + "Delta child decoded to a non fixed-width block", + )); + }; + if child.bits_per_value != bits_per_value || child.num_values != num_values - 1 { + return Err(Error::invalid_input(format!( + "Delta child decoded {} {}-bit values, expected {} {}-bit values", + child.num_values, + child.bits_per_value, + num_values - 1, + bits_per_value + ))); + } + + let data = match bits_per_value { + 32 => { + let deltas = checked_values::(&child, "Delta child")?; + let mut values = try_vec_with_capacity::(num_values, "Delta output")?; + let mut current = u32::try_from(base) + .map_err(|_| Error::invalid_input(format!("Delta base {base} exceeds u32::MAX")))?; + values.push(current); + for (index, delta) in deltas.iter().enumerate() { + current = current.checked_add(*delta).ok_or_else(|| { + Error::invalid_input(format!( + "Delta prefix sum overflows u32 at index {}", + index + 1 + )) + })?; + values.push(current); + } + LanceBuffer::reinterpret_vec(values) + } + 64 => { + let deltas = checked_values::(&child, "Delta child")?; + let mut values = try_vec_with_capacity::(num_values, "Delta output")?; + let mut current = base; + values.push(current); + for (index, delta) in deltas.iter().enumerate() { + current = current.checked_add(*delta).ok_or_else(|| { + Error::invalid_input(format!( + "Delta prefix sum overflows u64 at index {}", + index + 1 + )) + })?; + values.push(current); + } + LanceBuffer::reinterpret_vec(values) + } + _ => { + return Err(Error::invalid_input(format!( + "Delta only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + }; + + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value, + data, + num_values, + block_info: BlockInfo::default(), + })) +} + +fn checked_values( + data: &FixedWidthDataBlock, + label: &str, +) -> Result> { + let expected = usize::try_from(data.num_values) + .ok() + .and_then(|len| len.checked_mul(std::mem::size_of::())) + .ok_or_else(|| Error::invalid_input(format!("{label} byte length overflows usize")))?; + if data.data.len() != expected { + return Err(Error::invalid_input(format!( + "{label} has {} bytes, expected {expected}", + data.data.len() + ))); + } + Ok(data.data.borrow_to_typed_slice::()) +} + +#[cfg(test)] +mod tests { + use crate::compression::BlockCompressor; + use crate::encodings::physical::value::{ValueDecompressor, ValueEncoder}; + use crate::format::pb21::Flat; + + use super::*; + + #[test] + fn delta_round_trip_with_zero_delta() { + let input = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 64, + data: LanceBuffer::reinterpret_vec(vec![5_u64, 5, 9, 12]), + num_values: 4, + block_info: BlockInfo::default(), + }); + let DataBlock::FixedWidth(input) = input else { + unreachable!() + }; + let deltas = encode_deltas(input, 5).unwrap(); + let payload = ValueEncoder::default() + .compress(DataBlock::FixedWidth(deltas)) + .unwrap(); + let decoded = DeltaDecompressor::new( + 64, + 5, + Box::new(ValueDecompressor::from_flat(&Flat { + bits_per_value: 64, + data: None, + })), + ) + .decompress(payload, 4) + .unwrap() + .as_fixed_width() + .unwrap(); + assert_eq!( + decoded.data.borrow_to_typed_slice::().as_ref(), + &[5, 5, 9, 12] + ); + } + + #[test] + fn delta_rejects_decreasing_input() { + let input = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(vec![2_u32, 1]), + num_values: 2, + block_info: BlockInfo::default(), + }; + let error = encode_deltas(input, 2).unwrap_err(); + assert!(error.to_string().contains("decreases")); + } +} diff --git a/rust/lance-encoding/src/encodings/physical/range.rs b/rust/lance-encoding/src/encodings/physical/range.rs new file mode 100644 index 00000000000..23650c7e66c --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/range.rs @@ -0,0 +1,251 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! 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}, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, +}; +use lance_core::{Error, Result}; + +/// Returns the final value of an arithmetic sequence after checking its declared width. +pub(crate) fn checked_range_last( + bits_per_value: u64, + start: u64, + step: u64, + num_values: u64, +) -> Result { + if !matches!(bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Range only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + if step == 0 { + return Err(Error::invalid_input("Range step must be positive")); + } + if num_values < 2 { + return Err(Error::invalid_input(format!( + "Range requires at least 2 values, got {num_values}" + ))); + } + + let distance = step.checked_mul(num_values - 1).ok_or_else(|| { + Error::invalid_input(format!( + "Range step multiplication overflows: step={step}, num_values={num_values}" + )) + })?; + let last = start.checked_add(distance).ok_or_else(|| { + Error::invalid_input(format!( + "Range final value overflows: start={start}, step={step}, num_values={num_values}" + )) + })?; + if bits_per_value == 32 && last > u32::MAX as u64 { + return Err(Error::invalid_input(format!( + "Range final value {last} exceeds u32::MAX" + ))); + } + Ok(last) +} + +/// Validates an input block against the selected arithmetic range codec. +#[cfg(test)] +#[derive(Debug)] +pub(crate) struct RangeEncoder { + bits_per_value: u64, + start: u64, + step: u64, +} + +#[cfg(test)] +impl RangeEncoder { + pub(crate) fn new(bits_per_value: u64, start: u64, step: u64) -> Self { + Self { + bits_per_value, + start, + step, + } + } +} + +#[cfg(test)] +impl BlockCompressor for RangeEncoder { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Range encoding requires a fixed-width data block", + )); + }; + if data.bits_per_value != self.bits_per_value { + return Err(Error::invalid_input(format!( + "Range codec expects {}-bit values, got {}", + self.bits_per_value, data.bits_per_value + ))); + } + checked_range_last(self.bits_per_value, self.start, self.step, data.num_values)?; + + match self.bits_per_value { + 32 => { + let values = checked_values::(&data)?; + for (index, value) in values.iter().enumerate() { + let expected = self + .start + .checked_add(self.step.checked_mul(index as u64).ok_or_else(|| { + Error::invalid_input("Range index multiplication overflows") + })?) + .ok_or_else(|| Error::invalid_input("Range value addition overflows"))?; + if u64::from(*value) != expected { + return Err(Error::invalid_input(format!( + "Range input mismatch at index {index}: expected {expected}, got {value}" + ))); + } + } + } + 64 => { + let values = checked_values::(&data)?; + for (index, value) in values.iter().enumerate() { + let expected = self + .start + .checked_add(self.step.checked_mul(index as u64).ok_or_else(|| { + Error::invalid_input("Range index multiplication overflows") + })?) + .ok_or_else(|| Error::invalid_input("Range value addition overflows"))?; + if *value != expected { + return Err(Error::invalid_input(format!( + "Range input mismatch at index {index}: expected {expected}, got {value}" + ))); + } + } + } + _ => unreachable!("range width was validated above"), + } + Ok(None) + } +} + +/// Materializes a metadata-only arithmetic range. +#[derive(Debug)] +pub(crate) struct RangeDecompressor { + bits_per_value: u64, + start: u64, + step: u64, +} + +impl RangeDecompressor { + pub(crate) fn new(bits_per_value: u64, start: u64, step: u64) -> Self { + Self { + bits_per_value, + start, + step, + } + } +} + +impl BlockDecompressor for RangeDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + require_no_block_payload(data, "Range")?; + checked_range_last(self.bits_per_value, self.start, self.step, num_values)?; + materialize_validated_range(self.bits_per_value, self.start, self.step, num_values) + } +} + +pub(crate) fn materialize_validated_range( + bits_per_value: u64, + start: u64, + step: u64, + num_values: u64, +) -> Result { + let data = match bits_per_value { + 32 => { + let mut current = u32::try_from(start).map_err(|_| { + Error::invalid_input(format!("Range value {start} exceeds u32::MAX")) + })?; + let step = u32::try_from(step) + .map_err(|_| Error::invalid_input(format!("Range step {step} exceeds u32::MAX")))?; + let mut values = try_vec_with_capacity::(num_values, "Range output")?; + for index in 0..num_values { + values.push(current); + if index + 1 < num_values { + current += step; + } + } + LanceBuffer::reinterpret_vec(values) + } + 64 => { + let mut current = start; + let mut values = try_vec_with_capacity::(num_values, "Range output")?; + for index in 0..num_values { + values.push(current); + if index + 1 < num_values { + current += step; + } + } + LanceBuffer::reinterpret_vec(values) + } + _ => { + return Err(Error::invalid_input(format!( + "Range only supports 32 or 64-bit values, got {bits_per_value}" + ))); + } + }; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value, + data, + num_values, + block_info: BlockInfo::default(), + })) +} + +#[cfg(test)] +fn checked_values( + data: &FixedWidthDataBlock, +) -> Result> { + let expected = usize::try_from(data.num_values) + .ok() + .and_then(|len| len.checked_mul(std::mem::size_of::())) + .ok_or_else(|| Error::invalid_input("Range input byte length overflows usize"))?; + if data.data.len() != expected { + return Err(Error::invalid_input(format!( + "Range input has {} bytes, expected {expected}", + data.data.len() + ))); + } + Ok(data.data.borrow_to_typed_slice::()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn range_round_trip_u32() { + let input = DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(vec![3_u32, 8, 13, 18]), + num_values: 4, + block_info: BlockInfo::default(), + }); + let payload = RangeEncoder::new(32, 3, 5).compress(input).unwrap(); + assert!(payload.is_none()); + + let decoded = RangeDecompressor::new(32, 3, 5) + .decompress(payload, 4) + .unwrap() + .as_fixed_width() + .unwrap(); + assert_eq!( + decoded.data.borrow_to_typed_slice::().as_ref(), + &[3, 8, 13, 18] + ); + } + + #[test] + fn range_rejects_overflow() { + let error = checked_range_last(32, u32::MAX as u64, 1, 2).unwrap_err(); + assert!(error.to_string().contains("exceeds u32::MAX")); + } +} diff --git a/rust/lance-encoding/src/encodings/physical/rle.rs b/rust/lance-encoding/src/encodings/physical/rle.rs index 4091e0ddefb..b3d1f91757e 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -2006,6 +2006,7 @@ fn try_block_frame(values_payload_bytes: usize, lengths_payload_bytes: usize) -> #[derive(Debug, Clone, Copy)] pub(crate) enum MetadataRunLengths { Constant(u64), + Range { start: u64, step: u64 }, } impl MetadataRunLengths { @@ -2028,6 +2029,39 @@ impl MetadataRunLengths { } Ok(run_count) } + Self::Range { start, step } => { + if start == 0 { + return Err(Error::invalid_input("RLE run lengths must be positive")); + } + if step == 0 { + return Err(Error::invalid_input( + "RLE run lengths Range step must be positive", + )); + } + let target = u128::from(num_values); + let mut low = 1_u64; + let mut high = num_values; + while low <= high { + let run_count = low + (high - low) / 2; + let count = u128::from(run_count); + let factor = u128::from(start).checked_mul(2).and_then(|twice_start| { + u128::from(run_count - 1) + .checked_mul(u128::from(step)) + .and_then(|tail| twice_start.checked_add(tail)) + }); + let sum = factor + .and_then(|factor| count.checked_mul(factor)) + .map(|sum| sum / 2); + match sum.map(|sum| sum.cmp(&target)) { + Some(std::cmp::Ordering::Equal) => return Ok(run_count), + Some(std::cmp::Ordering::Less) => low = run_count + 1, + Some(std::cmp::Ordering::Greater) | None => high = run_count - 1, + } + } + Err(Error::invalid_input(format!( + "RLE run length range start={start} step={step} does not sum to {num_values}" + ))) + } } } } diff --git a/rust/lance-encoding/src/format.rs b/rust/lance-encoding/src/format.rs index f37e69b0216..74f1eafb831 100644 --- a/rust/lance-encoding/src/format.rs +++ b/rust/lance-encoding/src/format.rs @@ -670,6 +670,42 @@ macro_rules! impl_common_protobuf_utils { impl_common_protobuf_utils!(pb21, ProtobufUtils21); impl ProtobufUtils21 { + pub fn range( + uncompressed_bits_per_value: u64, + start: u64, + step: u64, + ) -> crate::format::pb21::CompressiveEncoding { + crate::format::pb21::CompressiveEncoding { + compression: Some( + crate::format::pb21::compressive_encoding::Compression::Range( + crate::format::pb21::Range { + uncompressed_bits_per_value, + start, + step, + }, + ), + ), + } + } + + pub fn delta( + uncompressed_bits_per_value: u64, + base: u64, + deltas: crate::format::pb21::CompressiveEncoding, + ) -> crate::format::pb21::CompressiveEncoding { + crate::format::pb21::CompressiveEncoding { + compression: Some( + crate::format::pb21::compressive_encoding::Compression::Delta(Box::new( + crate::format::pb21::Delta { + uncompressed_bits_per_value, + base, + deltas: Some(Box::new(deltas)), + }, + )), + ), + } + } + pub fn constant_layout( def_meaning: &[DefinitionInterpretation], inline_value: Option>, diff --git a/rust/lance-encoding/src/testing.rs b/rust/lance-encoding/src/testing.rs index 9825cb5949d..230044d0b70 100644 --- a/rust/lance-encoding/src/testing.rs +++ b/rust/lance-encoding/src/testing.rs @@ -545,6 +545,8 @@ fn tag(e: &Compression) -> &'static str { FixedSizeList(_) => "fixed_size_list", PackedStruct(_) => "packed_struct", VariablePackedStruct(_) => "variable_packed_struct", + Range(_) => "range", + Delta(_) => "delta", } } @@ -571,6 +573,7 @@ fn child(c: &Compression) -> Option> { Fsst(f) => f.values.as_ref().map(|b| vec![b.as_ref()]), ByteStreamSplit(b) => b.values.as_ref().map(|b| vec![b.as_ref()]), General(g) => g.values.as_ref().map(|b| vec![b.as_ref()]), + Delta(d) => d.deltas.as_ref().map(|b| vec![b.as_ref()]), Dictionary(d) => { let mut children = Vec::new(); if let Some(values) = d.items.as_ref() {