From 61fda29939d68768eaa2d50c1ed7c65163efc0aa Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 16:10:06 +0800 Subject: [PATCH 1/2] feat(encoding): add generic block compression --- docs/src/format/file/encoding.md | 82 +- protos/encodings_v2_1.proto | 40 +- rust/lance-encoding/benches/decoder.rs | 142 +- rust/lance-encoding/benches/encoder.rs | 83 +- rust/lance-encoding/src/compression.rs | 526 ++++-- rust/lance-encoding/src/compression/block.rs | 80 + .../src/compression/block/factory.rs | 670 +++++++ .../src/compression/block/fixed.rs | 122 ++ .../src/compression/block/tests.rs | 413 +++++ .../src/encodings/logical/primitive.rs | 325 +++- .../encodings/logical/primitive/miniblock.rs | 57 +- .../src/encodings/logical/primitive/sparse.rs | 166 +- .../logical/primitive/sparse/writer.rs | 72 +- rust/lance-encoding/src/encodings/physical.rs | 35 + .../src/encodings/physical/binary.rs | 1637 +++++++++++++++-- .../src/encodings/physical/binary/offsets.rs | 1449 +++++++++++++++ .../physical/binary/offsets/selector.rs | 233 +++ .../physical/binary/offsets/statistics.rs | 178 ++ .../src/encodings/physical/bitpacking.rs | 156 +- .../src/encodings/physical/block.rs | 296 ++- .../encodings/physical/byte_stream_split.rs | 20 +- .../src/encodings/physical/constant.rs | 140 +- .../src/encodings/physical/delta.rs | 276 +++ .../src/encodings/physical/dictionary.rs | 319 ++++ .../src/encodings/physical/fsst.rs | 10 +- .../src/encodings/physical/general.rs | 141 +- .../src/encodings/physical/packed.rs | 10 +- .../src/encodings/physical/range.rs | 245 +++ .../src/encodings/physical/rle.rs | 766 +++++++- .../src/encodings/physical/value.rs | 122 +- rust/lance-encoding/src/format.rs | 54 + rust/lance-encoding/src/testing.rs | 3 + .../tests/compression_strategy.rs | 81 + 33 files changed, 8319 insertions(+), 630 deletions(-) create mode 100644 rust/lance-encoding/src/compression/block.rs create mode 100644 rust/lance-encoding/src/compression/block/factory.rs create mode 100644 rust/lance-encoding/src/compression/block/fixed.rs create mode 100644 rust/lance-encoding/src/compression/block/tests.rs create mode 100644 rust/lance-encoding/src/encodings/physical/binary/offsets.rs create mode 100644 rust/lance-encoding/src/encodings/physical/binary/offsets/selector.rs create mode 100644 rust/lance-encoding/src/encodings/physical/binary/offsets/statistics.rs create mode 100644 rust/lance-encoding/src/encodings/physical/delta.rs create mode 100644 rust/lance-encoding/src/encodings/physical/dictionary.rs create mode 100644 rust/lance-encoding/src/encodings/physical/range.rs create mode 100644 rust/lance-encoding/tests/compression_strategy.rs diff --git a/docs/src/format/file/encoding.md b/docs/src/format/file/encoding.md index 1cb83d581e9..71b293ce06f 100644 --- a/docs/src/format/file/encoding.md +++ b/docs/src/format/file/encoding.md @@ -584,7 +584,10 @@ 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) | ❓ | ❓ | +| Constant | ✅ (2.1) | ❓ | ✅ (2.3, variable offsets) | +| Range | ✅ (2.3) | ❌ | ✅ (2.3, variable offsets) | +| Delta | ✅ (2.3) | ❌ | ✅ (2.3, variable offsets) | +| Dictionary | ✅ (2.3) | ❌ | ✅ (2.3, variable offsets) | | Bitpacking | ✅ (2.1) | ❓ | ✅ (2.1) | | Fsst | ❓ | ✅ (2.1) | ✅ (2.1) | | Rle | ✅ (2.2) | ❌ | ✅ (2.1) | @@ -594,6 +597,64 @@ 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 encode unsigned `u32` and `u64` sequences with a shared descriptor +contract. Direct codecs such as Flat, bitpacking, RLE, and Dictionary also support non-monotonic values; Range and +Delta require non-decreasing input. The containing layout supplies the value type and cardinality. A descriptor +constructs a concrete decoder tree whose nodes own their child codecs and framing validation. + +Writers select a concrete compressor from a bounded set: `Constant`, `Range`, `Flat`, bitpacking, RLE, +`Dictionary`, `Delta`, and general compression. Candidate costs include the protobuf descriptor, codec framing, +buffer entries, and alignment. Payload estimates are exact except for `General`, which extrapolates from a bounded +sample; only the selected compressor is invoked to materialize payloads. Writers use the following canonical +metadata-only codecs: + +- An empty sequence uses `Constant` with no scalar. +- A non-empty constant sequence uses `Constant` with one little-endian scalar. +- An arithmetic progression with a positive step uses `Range`. + +The first block-dictionary writer only emits Dictionary for `u64` sequences. +The first block-sequence grammar permits `General` only as the outer root around a `Flat` child; readers reject +inner, sibling, repeated, or non-`Flat` `General` transforms. + +`Range` stores the unsigned width, first value, and positive step. The value at index `i` is +`start + step * i`; readers reject overflow and widths other than 32 or 64 bits. + +```protobuf +%%% proto.message.Range %%% +``` + +`Delta` stores the first value inline. Its child represents the `n - 1` non-negative adjacent differences. Delta +has a payload exactly when its child has one. Zero differences are valid. Readers reconstruct the sequence with +checked prefix sums. + +```protobuf +%%% proto.message.Delta %%% +``` + +RLE and block Dictionary expose no outer payload when both children are metadata-only. Otherwise, they +combine their children into one framed outer payload: + +```text +RLE payload: + u64 values_payload_bytes + values payload + run-lengths payload + +Dictionary payload: + u64 indices_payload_bytes + u64 items_payload_bytes + indices payload + dictionary-items payload +``` + +The framed length for a metadata-only child must be zero. Readers validate child cardinalities, run-length sums, +dictionary index bounds, and all frame boundaries. + +Lance 2.0 through 2.2 writers do not emit `Range`, `Delta`, or block Dictionary. Their block selector order +and payload shapes remain unchanged. + ### Flat Flat compression is the uncompressed representation of fixed-width data. There is a single buffer of data @@ -611,6 +672,25 @@ When applied in a mini-block context each block may have a different number of v until we find the point that would exceed 4,096 bytes and then use the most recent power of 2 number of values that we have passed. +Lance 2.0 through 2.2 store each mini-block as one legacy buffer containing chunk-local Flat offsets followed by +the value bytes. Starting in Lance 2.3, the writer also evaluates a generic-offset container. Generic offsets are +zero-based and independently encoded in each chunk with one page-wide concrete offset codec: + +```text +Legacy chunk: + [adjusted Flat offsets][value bytes] + +Generic chunk: + [optional offset payload][value bytes] +``` + +The generic form has one buffer size per chunk for metadata-only offset codecs and two for payload-bearing offset +codecs. It is selected only when its complete serialized size is strictly smaller than the legacy form. A Flat +offset descriptor always denotes the legacy form, which keeps the two wire shapes unambiguous. Both forms keep +offsets and values in the same mini-block chunk, so a random read still fetches only the selected chunk. Fields that +explicitly request an outer General compressor retain the legacy form because that transform can change the +complete-container winner after offset selection. + ### Constant Constant compression is currently only utilized in a few specialized scenarios such as all-null arrays. diff --git a/protos/encodings_v2_1.proto b/protos/encodings_v2_1.proto index 51427332063..1a7c0dbcaf2 100644 --- a/protos/encodings_v2_1.proto +++ b/protos/encodings_v2_1.proto @@ -406,8 +406,9 @@ message Flat { // This is a transparent encoding by definition. // // The input is a variable-width data block. -// The output is a single fixed-width data block (the offsets) and -// a single buffer (the values) +// The output is an optional encoded offset payload and one values buffer. +// Legacy mini-block Variable encoding instead stores Flat offsets and values +// interleaved in one buffer. message Variable { // Describes how the offsets data block is compressed CompressiveEncoding offsets = 1; @@ -426,6 +427,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 +665,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/benches/decoder.rs b/rust/lance-encoding/benches/decoder.rs index cc0404e1bb3..f4aecf69522 100644 --- a/rust/lance-encoding/benches/decoder.rs +++ b/rust/lance-encoding/benches/decoder.rs @@ -2,7 +2,7 @@ // SPDX-FileCopyrightText: Copyright The Lance Authors use std::{collections::HashMap, sync::Arc}; -use arrow_array::{RecordBatch, UInt32Array}; +use arrow_array::{RecordBatch, StringArray, UInt32Array}; use arrow_schema::{DataType, Field, Schema, TimeUnit}; use arrow_select::take::take; use criterion::{Criterion, criterion_group, criterion_main}; @@ -10,6 +10,7 @@ use futures::StreamExt; use lance_core::cache::LanceCache; use lance_datagen::ArrayGeneratorExt; use lance_encoding::{ + BufferScheduler, EncodingsIo, decoder::{ DecodeBatchScheduler, DecoderConfig, DecoderPlugins, FilterExpression, create_decode_stream, }, @@ -563,6 +564,140 @@ fn bench_decode_compressed_parallel(c: &mut Criterion) { } } +async fn decode_take( + encoded: &lance_encoding::encoder::EncodedBatch, + indices: &[u64], + cache: Arc, +) -> RecordBatch { + let io_scheduler = Arc::new(BufferScheduler::new(encoded.data.clone())) as Arc; + let filter = FilterExpression::no_filter(); + let mut decode_scheduler = DecodeBatchScheduler::try_new( + encoded.schema.as_ref(), + &encoded.top_level_columns, + &encoded.page_table, + &vec![], + encoded.num_rows, + Arc::::default(), + io_scheduler.clone(), + cache, + &filter, + &DecoderConfig::default(), + ) + .await + .unwrap(); + let (tx, rx) = unbounded_channel(); + decode_scheduler.schedule_take(indices, &filter, tx, io_scheduler); + let mut stream = create_decode_stream( + &encoded.schema, + indices.len() as u64, + indices.len() as u32, + true, + false, + true, + rx, + None, + ) + .unwrap(); + stream.next().await.unwrap().task.await.unwrap() +} + +fn bench_variable_offsets_decode(c: &mut Criterion) { + const NUM_ROWS: usize = 262_144; + const NUM_TAKES: usize = 512; + + let metadata = HashMap::from([ + ( + "lance-encoding:structural-encoding".to_string(), + "miniblock".to_string(), + ), + ( + "lance-encoding:dict-divisor".to_string(), + "100000".to_string(), + ), + ("lance-encoding:compression".to_string(), "none".to_string()), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Utf8, false).with_metadata(metadata), + ])); + let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(schema.as_ref()).unwrap()); + let corpora = [ + ( + "range", + Arc::new(StringArray::from_iter_values( + (0..NUM_ROWS).map(|index| format!("row_{index:012}")), + )) as Arc, + ), + ( + "delta", + Arc::new(StringArray::from_iter_values( + (0..NUM_ROWS).map(|index| "x".repeat(4 + index % 64)), + )) as Arc, + ), + ]; + let mut take_indices = (0..NUM_TAKES) + .map(|index| (index as u64).wrapping_mul(104_729).wrapping_add(8_191) % NUM_ROWS as u64) + .collect::>(); + take_indices.sort_unstable(); + take_indices.dedup(); + + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("variable_offsets_decode"); + + for (workload, array) in corpora { + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + for version in [LanceFileVersion::V2_2, LanceFileVersion::V2_3] { + let strategy = default_encoding_strategy(version); + let options = EncodingOptions { + version, + ..Default::default() + }; + let encoded = rt + .block_on(encode_batch( + &batch, + lance_schema.clone(), + strategy.as_ref(), + &options, + )) + .unwrap(); + let encoded_bytes = encoded.data.len(); + let benchmark_suffix = format!("{workload}/{version}/{encoded_bytes}B"); + + group.throughput(criterion::Throughput::Elements(NUM_ROWS as u64)); + group.bench_function(format!("scan/{benchmark_suffix}"), |bencher| { + bencher.iter(|| { + let decoded = rt + .block_on(lance_encoding::decoder::decode_batch( + &encoded, + &FilterExpression::no_filter(), + Arc::::default(), + false, + version, + Some(Arc::new(LanceCache::no_cache())), + )) + .unwrap(); + assert_eq!(decoded.num_rows(), NUM_ROWS); + }) + }); + + for cache_mode in ["cold", "warm"] { + let cache = if cache_mode == "cold" { + Arc::new(LanceCache::no_cache()) + } else { + Arc::new(LanceCache::with_capacity(64 * 1024 * 1024)) + }; + group.throughput(criterion::Throughput::Elements(take_indices.len() as u64)); + group.bench_function(format!("take/{benchmark_suffix}/{cache_mode}"), |bencher| { + bencher.iter(|| { + let decoded = + rt.block_on(decode_take(&encoded, &take_indices, cache.clone())); + assert_eq!(decoded.num_rows(), take_indices.len()); + }) + }); + } + } + } +} + #[cfg(target_os = "linux")] criterion_group!( name=benches; @@ -570,7 +705,7 @@ criterion_group!( .with_profiler(lance_testing::pprof::PProfProfiler::new(100, lance_testing::pprof::Output::Flamegraph(None))); targets = bench_decode, bench_decode_fsl, bench_decode_str_with_dict_encoding, bench_decode_packed_struct, bench_decode_str_with_fixed_size_binary_encoding, bench_decode_compressed, - bench_decode_compressed_parallel); + bench_decode_compressed_parallel, bench_variable_offsets_decode); // Non-linux version does not support pprof. #[cfg(not(target_os = "linux"))] @@ -578,5 +713,6 @@ criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10); targets = bench_decode, bench_decode_fsl, bench_decode_str_with_dict_encoding, bench_decode_packed_struct, - bench_decode_compressed, bench_decode_compressed_parallel); + bench_decode_compressed, bench_decode_compressed_parallel, + bench_variable_offsets_decode); criterion_main!(benches); diff --git a/rust/lance-encoding/benches/encoder.rs b/rust/lance-encoding/benches/encoder.rs index 08eb89d32fe..e766a557339 100644 --- a/rust/lance-encoding/benches/encoder.rs +++ b/rust/lance-encoding/benches/encoder.rs @@ -1,12 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, hint::black_box, sync::Arc}; -use arrow_array::{ArrayRef, BooleanArray, ListArray, RecordBatch}; +use arrow_array::{ArrayRef, BooleanArray, ListArray, RecordBatch, StringArray}; use arrow_buffer::{OffsetBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field, Schema}; -use criterion::{Criterion, criterion_group, criterion_main}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; use lance_encoding::{ encoder::{EncodingOptions, default_encoding_strategy, encode_batch}, version::LanceFileVersion, @@ -162,17 +162,90 @@ fn bench_encode_structural_pages(c: &mut Criterion) { }); } +fn bench_variable_offsets(c: &mut Criterion) { + const NUM_ROWS: usize = 262_144; + let metadata = HashMap::from([ + ( + "lance-encoding:structural-encoding".to_string(), + "miniblock".to_string(), + ), + ( + "lance-encoding:dict-divisor".to_string(), + "100000".to_string(), + ), + ("lance-encoding:compression".to_string(), "none".to_string()), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Utf8, false).with_metadata(metadata), + ])); + let lance_schema = Arc::new(lance_core::datatypes::Schema::try_from(schema.as_ref()).unwrap()); + let runtime = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("variable_offsets"); + group.throughput(Throughput::Elements(NUM_ROWS as u64)); + let corpora: [(&str, ArrayRef); 2] = [ + ( + "range", + Arc::new(StringArray::from_iter_values( + (0..NUM_ROWS).map(|index| format!("row_{index:012}")), + )), + ), + ( + "delta", + Arc::new(StringArray::from_iter_values( + (0..NUM_ROWS).map(|index| "x".repeat(4 + index % 64)), + )), + ), + ]; + for (workload, array) in corpora { + let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + for version in [LanceFileVersion::V2_2, LanceFileVersion::V2_3] { + let strategy = default_encoding_strategy(version); + let options = EncodingOptions { + version, + ..Default::default() + }; + let encoded_bytes = runtime + .block_on(encode_batch( + &batch, + lance_schema.clone(), + strategy.as_ref(), + &options, + )) + .unwrap() + .data + .len(); + group.bench_function( + format!("{workload}/{version}/{encoded_bytes}B"), + |bencher| { + bencher.iter(|| { + black_box( + runtime + .block_on(encode_batch( + &batch, + lance_schema.clone(), + strategy.as_ref(), + &options, + )) + .unwrap(), + ) + }) + }, + ); + } + } +} + #[cfg(target_os = "linux")] criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10) .with_profiler(lance_testing::pprof::PProfProfiler::new(100, lance_testing::pprof::Output::Flamegraph(None))); - targets = bench_encode_compressed, bench_encode_structural_pages); + targets = bench_encode_compressed, bench_encode_structural_pages, bench_variable_offsets); #[cfg(not(target_os = "linux"))] criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10); - targets = bench_encode_compressed, bench_encode_structural_pages); + targets = bench_encode_compressed, bench_encode_structural_pages, bench_variable_offsets); criterion_main!(benches); diff --git a/rust/lance-encoding/src/compression.rs b/rust/lance-encoding/src/compression.rs index 5c15d48c25a..65d349b9d1d 100644 --- a/rust/lance-encoding/src/compression.rs +++ b/rust/lance-encoding/src/compression.rs @@ -16,6 +16,10 @@ //! Fullzip compression is a per-value approach where we require that values are transparently //! compressed so that we can locate them later. +pub mod block; + +pub use block::BlockValueType; + #[cfg(feature = "bitpacking")] use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking}; use crate::{ @@ -69,7 +73,9 @@ use crate::{ version::LanceFileVersion, }; -use arrow_array::{cast::AsArray, types::UInt64Type}; +#[cfg(feature = "bitpacking")] +use arrow_array::cast::AsArray; +use arrow_array::types::UInt64Type; use arrow_schema::DataType; use fsst::fsst::{FSST_LEAST_INPUT_MAX_LENGTH, FSST_LEAST_INPUT_SIZE}; use lance_core::{Error, Result, datatypes::Field, error::LanceOptionExt}; @@ -99,11 +105,11 @@ const RLE_BLOCK_HEADER_BYTES: u128 = std::mem::size_of::() as u128; /// required (e.g. when encoding metadata buffers like a dictionary or for encoding rep/def /// mini-block chunks) pub trait BlockCompressor: std::fmt::Debug + Send + Sync { - /// Compress the data into a single buffer + /// Compress the data into zero or one buffers. /// - /// Also returns a description of the compression that can be used to decompress - /// when reading the data back - fn compress(&self, data: DataBlock) -> Result; + /// `None` represents a metadata-only codec. `Some` represents a physical + /// payload, including a zero-byte payload. + fn compress(&self, data: DataBlock) -> Result>; } /// A trait to pick which compression to use for given data @@ -119,7 +125,10 @@ pub trait BlockCompressor: std::fmt::Debug + Send + Sync { /// used for narrow data types (both fixed and variable length) where we can /// fit many values into an 16KiB block. pub trait CompressionStrategy: Send + Sync + std::fmt::Debug { - /// Create a block compressor for the given data + /// Create a block compressor for the given data. + /// + /// The returned concrete codec may be reused for independently framed + /// blocks that use the same encoding. fn create_block_compressor( &self, field: &Field, @@ -141,6 +150,18 @@ pub trait CompressionStrategy: Send + Sync + std::fmt::Debug { ) -> Result>; } +pub(crate) fn compress_required_block( + strategy: &dyn CompressionStrategy, + field: &Field, + data: DataBlock, +) -> Result<(LanceBuffer, CompressiveEncoding)> { + let (compressor, encoding) = strategy.create_block_compressor(field, &data)?; + let payload = compressor.compress(data)?.ok_or_else(|| { + Error::internal("Required block compressor selected a metadata-only codec".to_string()) + })?; + Ok((payload, encoding)) +} + #[derive(Debug, Default, Clone)] pub struct DefaultCompressionStrategy { /// User-configured compression parameters @@ -370,8 +391,6 @@ fn try_bitpack_for_mini_block(_data: &FixedWidthDataBlock) -> Option Option { - use arrow_array::cast::AsArray; - let bits = data.bits_per_value; if !matches!(bits, 8 | 16 | 32 | 64) { return None; @@ -402,6 +421,7 @@ fn estimate_inline_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option { u64::try_from(estimated_bytes).ok() } +#[cfg(feature = "bitpacking")] fn try_bitpack_for_block( data: &FixedWidthDataBlock, ) -> Option<(Box, CompressiveEncoding)> { @@ -435,6 +455,13 @@ fn try_bitpack_for_block( } } +#[cfg(not(feature = "bitpacking"))] +fn try_bitpack_for_block( + _data: &FixedWidthDataBlock, +) -> Option<(Box, CompressiveEncoding)> { + None +} + #[cfg(feature = "bitpacking")] fn estimate_block_bitpacking_bytes(data: &FixedWidthDataBlock) -> Option { let bits = data.bits_per_value; @@ -657,9 +684,21 @@ impl DefaultCompressionStrategy { let data_size = data.expect_single_stat::(Stat::DataSize); let max_len = data.expect_single_stat::(Stat::MaxLength); - // Explicitly disable all compression. + let build_binary = || { + let generic_offsets_own_final_payload = + compression.is_none() || compression == Some("none"); + if self.version.resolve() >= LanceFileVersion::V2_3 && generic_offsets_own_final_payload + { + BinaryMiniBlockEncoder::with_generic_offsets(params.minichunk_size, params.clone()) + } else { + BinaryMiniBlockEncoder::new(params.minichunk_size) + } + }; + + // "none" disables general compression but still permits structural + // offset codecs in 2.3. if compression == Some("none") { - return Ok(Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size))); + return Ok(Box::new(build_binary())); } let use_fsst = compression == Some("fsst") @@ -672,7 +711,7 @@ impl DefaultCompressionStrategy { let mut base_encoder: Box = if use_fsst { Box::new(FsstMiniBlockEncoder::new(params.minichunk_size)) } else { - Box::new(BinaryMiniBlockEncoder::new(params.minichunk_size)) + Box::new(build_binary()) }; // Wrap with general compression when configured (except FSST / none). @@ -701,6 +740,63 @@ impl DefaultCompressionStrategy { } impl CompressionStrategy for DefaultCompressionStrategy { + fn create_block_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result<(Box, CompressiveEncoding)> { + let field_params = self.get_merged_field_params(field); + + match data { + DataBlock::FixedWidth(fixed_width) => { + if let Some((compressor, encoding)) = + try_rle_for_block(fixed_width, self.version, &field_params, self.use_rle_v2())? + { + return Ok((compressor, encoding)); + } + if let Some((compressor, encoding)) = try_bitpack_for_block(fixed_width) { + return Ok((compressor, encoding)); + } + + if let Some((compressor, config)) = + try_general_compression(self.version, &field_params, data)? + { + let encoding = ProtobufUtils21::wrapped( + config, + ProtobufUtils21::flat(fixed_width.bits_per_value, None), + )?; + return Ok((compressor, encoding)); + } + + let encoder = Box::new(ValueEncoder::default()); + let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); + Ok((encoder, encoding)) + } + DataBlock::VariableWidth(variable_width) => { + if let Some((compressor, config)) = + try_general_compression(self.version, &field_params, data)? + { + let encoding = ProtobufUtils21::wrapped( + config, + ProtobufUtils21::variable( + ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), + None, + ), + )?; + return Ok((compressor, encoding)); + } + + let encoder = Box::new(VariableEncoder::default()); + let encoding = ProtobufUtils21::variable( + ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), + None, + ); + Ok((encoder, encoding)) + } + _ => unreachable!(), + } + } + fn create_miniblock_compressor( &self, field: &Field, @@ -835,65 +931,6 @@ impl CompressionStrategy for DefaultCompressionStrategy { ), } } - - fn create_block_compressor( - &self, - field: &Field, - data: &DataBlock, - ) -> Result<(Box, CompressiveEncoding)> { - let field_params = self.get_merged_field_params(field); - - match data { - DataBlock::FixedWidth(fixed_width) => { - if let Some((compressor, encoding)) = - try_rle_for_block(fixed_width, self.version, &field_params, self.use_rle_v2())? - { - return Ok((compressor, encoding)); - } - if let Some((compressor, encoding)) = try_bitpack_for_block(fixed_width) { - return Ok((compressor, encoding)); - } - - // Try general compression (user-requested or automatic over MIN_BLOCK_SIZE_FOR_GENERAL_COMPRESSION) - if let Some((compressor, config)) = - try_general_compression(self.version, &field_params, data)? - { - let encoding = ProtobufUtils21::wrapped( - config, - ProtobufUtils21::flat(fixed_width.bits_per_value, None), - )?; - return Ok((compressor, encoding)); - } - - let encoder = Box::new(ValueEncoder::default()); - let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); - Ok((encoder, encoding)) - } - DataBlock::VariableWidth(variable_width) => { - // Try general compression - if let Some((compressor, config)) = - try_general_compression(self.version, &field_params, data)? - { - let encoding = ProtobufUtils21::wrapped( - config, - ProtobufUtils21::variable( - ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), - None, - ), - )?; - return Ok((compressor, encoding)); - } - - let encoder = Box::new(VariableEncoder::default()); - let encoding = ProtobufUtils21::variable( - ProtobufUtils21::flat(variable_width.bits_per_offset as u64, None), - None, - ); - Ok((encoder, encoding)) - } - _ => unreachable!(), - } - } } pub trait MiniBlockDecompressor: std::fmt::Debug + Send + Sync { @@ -915,7 +952,18 @@ pub trait VariablePerValueDecompressor: std::fmt::Debug + Send + Sync { } pub trait BlockDecompressor: std::fmt::Debug + Send + Sync { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result; + fn decompress(&self, data: Option, num_values: u64) -> Result; +} + +pub(crate) fn require_block_payload(data: Option, codec: &str) -> Result { + data.ok_or_else(|| Error::invalid_input(format!("{codec} requires one payload"))) +} + +pub(crate) fn require_no_block_payload(data: Option, codec: &str) -> Result<()> { + if data.is_some() { + return Err(Error::invalid_input(format!("{codec} expects no payload"))); + } + Ok(()) } pub trait DecompressionStrategy: std::fmt::Debug + Send + Sync { @@ -950,7 +998,10 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { description: &CompressiveEncoding, decompression_strategy: &dyn DecompressionStrategy, ) -> Result> { - match description.compression.as_ref().unwrap() { + let compression = description.compression.as_ref().ok_or_else(|| { + Error::invalid_input("Mini-block encoding is missing its compression variant") + })?; + match compression { Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))), #[cfg(feature = "bitpacking")] Compression::InlineBitpacking(description) => { @@ -960,24 +1011,14 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Compression::InlineBitpacking(_) => Err(Error::not_supported_source( "this runtime was not built with bitpacking support".into(), )), - Compression::Variable(variable) => { - let Compression::Flat(offsets) = variable - .offsets - .as_ref() - .unwrap() - .compression - .as_ref() - .unwrap() - else { - panic!("Variable compression only supports flat offsets") - }; - Ok(Box::new(BinaryMiniBlockDecompressor::new( - offsets.bits_per_value as u8, - ))) - } + Compression::Variable(variable) => Ok(Box::new( + BinaryMiniBlockDecompressor::from_variable(variable)?, + )), Compression::Fsst(description) => { let inner_decompressor = decompression_strategy.create_miniblock_decompressor( - description.values.as_ref().unwrap(), + description.values.as_ref().ok_or_else(|| { + Error::invalid_input("FSST mini-block is missing its values encoding") + })?, decompression_strategy, )?; Ok(Box::new(FsstMiniBlockDecompressor::new( @@ -1001,10 +1042,18 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { decompression_strategy, )?)), Compression::ByteStreamSplit(bss) => { - let Compression::Flat(values) = - bss.values.as_ref().unwrap().compression.as_ref().unwrap() + let values = bss.values.as_ref().ok_or_else(|| { + Error::invalid_input("ByteStreamSplit is missing its values encoding") + })?; + let Compression::Flat(values) = values.compression.as_ref().ok_or_else(|| { + Error::invalid_input( + "ByteStreamSplit values are missing their compression variant", + ) + })? else { - panic!("ByteStreamSplit compression only supports flat values") + return Err(Error::invalid_input( + "ByteStreamSplit compression only supports flat values", + )); }; Ok(Box::new(ByteStreamSplitDecompressor::new( values.bits_per_value as usize, @@ -1033,7 +1082,13 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { compression_config, ))) } - _ => todo!(), + other => Err(Error::not_supported_source( + format!( + "Mini-block decompression does not support {} encoding", + compression_name(other) + ) + .into(), + )), } } @@ -1128,44 +1183,70 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { &self, description: &CompressiveEncoding, ) -> Result> { - match description.compression.as_ref().unwrap() { - Compression::InlineBitpacking(inline_bitpacking) => Ok(Box::new( - InlineBitpacking::from_description(inline_bitpacking), - )), + let compression = description.compression.as_ref().ok_or_else(|| { + Error::invalid_input("Block encoding is missing its compression variant") + })?; + match compression { Compression::Flat(flat) => Ok(Box::new(ValueDecompressor::from_flat(flat))), - Compression::Constant(constant) => { - let scalar = constant + Compression::Constant(constant) => Ok(Box::new(ConstantDecompressor::new( + constant .value .as_ref() - .map(|v| LanceBuffer::from_bytes(v.clone(), 1)); - Ok(Box::new(ConstantDecompressor::new(scalar))) - } - Compression::Variable(_) => Ok(Box::new(BinaryBlockDecompressor::default())), - Compression::FixedSizeList(fsl) => { - Ok(Box::new(ValueDecompressor::from_fsl(fsl.as_ref()))) + .map(|value| LanceBuffer::from_bytes(value.clone(), 1)), + ))), + #[cfg(feature = "bitpacking")] + Compression::InlineBitpacking(inline) => { + Ok(Box::new(InlineBitpacking::from_description(inline))) } + #[cfg(not(feature = "bitpacking"))] + Compression::InlineBitpacking(_) => Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )), + #[cfg(feature = "bitpacking")] Compression::OutOfLineBitpacking(out_of_line) => { - // Extract the compressed bit width from the values encoding - let compressed_bit_width = match out_of_line - .values - .as_ref() - .unwrap() - .compression - .as_ref() - .unwrap() - { - Compression::Flat(flat) => flat.bits_per_value, - _ => { - return Err(Error::invalid_input_source( - "OutOfLineBitpacking values must use Flat encoding".into(), - )); - } + let values = out_of_line.values.as_deref().ok_or_else(|| { + Error::invalid_input("OutOfLineBitpacking is missing its values encoding") + })?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input( + "OutOfLineBitpacking values must use Flat encoding", + )); }; Ok(Box::new(OutOfLineBitpacking::new( - compressed_bit_width, + flat.bits_per_value, out_of_line.uncompressed_bits_per_value, ))) } + #[cfg(not(feature = "bitpacking"))] + Compression::OutOfLineBitpacking(_) => Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )), + Compression::Dictionary(_) | 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(|| { + Error::invalid_input("Variable block encoding is missing offsets") + })?; + let Some(Compression::Flat(flat)) = offsets.compression.as_ref() else { + return Err(Error::invalid_input( + "Variable block encoding only supports flat offsets", + )); + }; + if !matches!(flat.bits_per_value, 32 | 64) || flat.data.is_some() { + return Err(Error::invalid_input(format!( + "Variable block offsets require uncompressed 32 or 64-bit Flat encoding, got {} bits", + flat.bits_per_value + ))); + } + Ok(Box::new(BinaryBlockDecompressor::default())) + } + Compression::FixedSizeList(fsl) => { + Ok(Box::new(ValueDecompressor::from_fsl(fsl.as_ref()))) + } Compression::General(general) => { let inner_desc = general .values @@ -1186,8 +1267,13 @@ impl DecompressionStrategy for DefaultDecompressionStrategy { Ok(Box::new(general_decompressor)) } - Compression::Rle(rle) => Ok(Box::new(create_rle_decompressor(rle, self)?)), - _ => todo!(), + other => Err(Error::not_supported_source( + format!( + "Block decompression does not support {} encoding", + compression_name(other) + ) + .into(), + )), } } } @@ -1256,6 +1342,21 @@ fn create_rle_child_decompressor( let (bits_per_value, requires_num_values, needs_decompressor) = validate_rle_child_compression(compression, role)?; + if let Compression::General(general) = compression + && !requires_num_values + { + let compression = general.compression.as_ref().ok_or_else(|| { + Error::invalid_input(format!( + "RLE {role} general child missing compression config" + )) + })?; + let scheme = compression.scheme().try_into()?; + return RleChildDecompressor::general( + bits_per_value, + CompressionConfig::new(scheme, compression.level), + ); + } + if needs_decompressor { Ok(RleChildDecompressor::block( bits_per_value, @@ -1359,6 +1460,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", } } @@ -1379,6 +1482,21 @@ mod tests { field } + fn selected_block_codec( + strategy: &dyn CompressionStrategy, + field: &Field, + data: &DataBlock, + ) -> (Box, CompressiveEncoding) { + strategy.create_block_compressor(field, data).unwrap() + } + + fn miniblock_context() + -> crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext { + crate::encodings::logical::primitive::miniblock::MiniBlockCompressionContext::new( + 0, true, true, + ) + } + fn create_fixed_width_block_with_stats( bits_per_value: u64, num_values: u64, @@ -1465,6 +1583,7 @@ mod tests { run_lengths.bits_per_value } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn expect_rle_encoding(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { match encoding.compression.as_ref().unwrap() { Compression::Rle(rle) => rle, @@ -1622,8 +1741,8 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let (compressor, _encoding) = strategy.create_block_compressor(&field, &data).unwrap(); - let debug_str = format!("{:?}", compressor); + let (compressor, _) = selected_block_codec(&strategy, &field, &data); + let debug_str = format!("{compressor:?}"); assert!( debug_str.contains("OutOfLineBitpacking"), "expected OutOfLineBitpacking, got: {debug_str}" @@ -1644,7 +1763,7 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); assert!(format!("{compressor:?}").contains("ValueEncoder")); assert!(matches!( @@ -1672,7 +1791,7 @@ mod tests { block.compute_stat(); let data = DataBlock::FixedWidth(block); - let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); let debug_str = format!("{compressor:?}"); assert!( debug_str.contains("OutOfLineBitpacking"), @@ -1759,12 +1878,16 @@ mod tests { let compressor = strategy .create_miniblock_compressor(&field, &fixed_data) .unwrap(); - let (_block, encoding) = compressor.compress(fixed_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(fixed_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, false); let compressor = strategy .create_miniblock_compressor(&field, &variable_data) .unwrap(); - let (_block, encoding) = compressor.compress(variable_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(variable_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, true); // Test pervalue @@ -1776,6 +1899,109 @@ mod tests { check_uncompressed_encoding(&encoding, true); } + #[test] + fn test_variable_offset_codec_is_version_gated() { + let num_values = 2_048_u64; + let offsets = (0..=num_values) + .map(|index| (index * 3) as i32) + .collect::>(); + let mut variable = VariableWidthBlock { + data: LanceBuffer::from(vec![7_u8; num_values as usize * 3]), + offsets: LanceBuffer::reinterpret_vec(offsets), + bits_per_offset: 32, + num_values, + block_info: BlockInfo::default(), + }; + variable.compute_stat(); + let data = DataBlock::VariableWidth(variable); + let field = create_test_field("bytes", DataType::Binary); + + for version in [ + LanceFileVersion::V2_0, + LanceFileVersion::V2_1, + LanceFileVersion::V2_2, + ] { + let strategy = DefaultCompressionStrategy::new().with_version(version); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (compressed, encoding) = compressor + .compress(data.clone(), miniblock_context()) + .unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding for {version}"); + }; + 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) + ); + } + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let Some(Compression::Variable(variable)) = encoding.compression.as_ref() else { + panic!("expected Variable encoding for 2.3"); + }; + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(Compression::Range(_)) + )); + assert_eq!(compressed.data.len(), 1); + } + + #[test] + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn test_v2_3_explicit_general_compression_keeps_legacy_offsets() { + let num_values = 2_048_u64; + let offsets = (0..=num_values) + .map(|index| (index * 3) as i32) + .collect::>(); + let mut variable = VariableWidthBlock { + data: LanceBuffer::from(vec![7_u8; num_values as usize * 3]), + offsets: LanceBuffer::reinterpret_vec(offsets), + bits_per_offset: 32, + num_values, + block_info: BlockInfo::default(), + }; + variable.compute_stat(); + let data = DataBlock::VariableWidth(variable); + let mut field = create_test_field("bytes", DataType::Binary); + field.metadata.insert( + COMPRESSION_META_KEY.to_string(), + if cfg!(feature = "lz4") { "lz4" } else { "zstd" }.to_string(), + ); + + let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); + let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); + let (_, encoding) = compressor.compress(data, miniblock_context()).unwrap(); + let value_encoding = match encoding.compression.as_ref().unwrap() { + Compression::General(general) => general.values.as_deref().unwrap(), + _ => &encoding, + }; + let Some(Compression::Variable(variable)) = value_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(_)) + )); + } + #[test] fn test_field_metadata_none_compression() { // Prepare field with metadata for none compression @@ -1794,13 +2020,17 @@ mod tests { let compressor = strategy .create_miniblock_compressor(&field, &fixed_data) .unwrap(); - let (_block, encoding) = compressor.compress(fixed_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(fixed_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, false); let compressor = strategy .create_miniblock_compressor(&field, &variable_data) .unwrap(); - let (_block, encoding) = compressor.compress(variable_data.clone()).unwrap(); + let (_block, encoding) = compressor + .compress(variable_data.clone(), miniblock_context()) + .unwrap(); check_uncompressed_encoding(&encoding, true); // Test pervalue @@ -2097,7 +2327,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2122,7 +2352,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(version); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8, "version={version}"); } } @@ -2150,7 +2380,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2173,7 +2403,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 16); } @@ -2196,7 +2426,7 @@ mod tests { let strategy = DefaultCompressionStrategy::new().with_version(LanceFileVersion::V2_3); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); assert_eq!(rle_run_length_bits(&encoding), 8); } @@ -2233,7 +2463,7 @@ mod tests { let data = DataBlock::FixedWidth(data); let compressor = strategy.create_miniblock_compressor(&field, &data).unwrap(); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let rle = expect_rle_encoding(&encoding); assert!( @@ -2281,7 +2511,7 @@ mod tests { let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let Compression::Rle(rle) = encoding.compression.as_ref().unwrap() else { panic!("expected RLE encoding"); }; @@ -2331,7 +2561,7 @@ mod tests { "expected RLE to beat inline bitpacking after child selection, got: {debug_str}" ); - let (_compressed, encoding) = compressor.compress(data).unwrap(); + let (_compressed, encoding) = compressor.compress(data, miniblock_context()).unwrap(); let rle = expect_rle_encoding(&encoding); assert!(matches!( rle.values.as_ref().unwrap().compression.as_ref().unwrap(), @@ -2523,7 +2753,7 @@ mod tests { let field = create_test_field("dict_values", DataType::FixedSizeBinary(3)); let data = create_fixed_width_block(24, 1024); - let (_compressor, encoding) = strategy + let (_, encoding) = strategy .create_block_compressor(&field, &data) .expect("block compressor selection should succeed"); @@ -2554,7 +2784,7 @@ mod tests { "test requires block size above automatic general compression threshold" ); - let (_compressor, encoding) = strategy + let (_, encoding) = strategy .create_block_compressor(&field, &data) .expect("block compressor selection should succeed"); @@ -2579,7 +2809,7 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_3); - let (compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (compressor, encoding) = selected_block_codec(&strategy, &field, &data); assert_eq!(rle_run_length_bits(&encoding), 32); let compressed = compressor.compress(data).unwrap(); @@ -2614,7 +2844,7 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_2); - let (_compressor, encoding) = strategy.create_block_compressor(&field, &data).unwrap(); + let (_, encoding) = selected_block_codec(&strategy, &field, &data); assert_eq!(rle_run_length_bits(&encoding), 8); } @@ -2645,11 +2875,9 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_2); - let (compressor, _) = strategy - .create_block_compressor(&field, &data_block) - .unwrap(); + let (compressor, _) = selected_block_codec(&strategy, &field, &data_block); - let debug_str = format!("{:?}", compressor); + let debug_str = format!("{compressor:?}"); assert!(debug_str.contains("RleEncoder")); } @@ -2680,11 +2908,9 @@ mod tests { let strategy = DefaultCompressionStrategy::with_params(CompressionParams::new()) .with_version(LanceFileVersion::V2_1); - let (compressor, _) = strategy - .create_block_compressor(&field, &data_block) - .unwrap(); + let (compressor, _) = selected_block_codec(&strategy, &field, &data_block); - let debug_str = format!("{:?}", compressor); + let debug_str = format!("{compressor:?}"); assert!( !debug_str.contains("RleEncoder"), "RLE should not be used for V2.1" diff --git a/rust/lance-encoding/src/compression/block.rs b/rust/lance-encoding/src/compression/block.rs new file mode 100644 index 00000000000..8974bc866d6 --- /dev/null +++ b/rust/lance-encoding/src/compression/block.rs @@ -0,0 +1,80 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Shared types and decoder construction for unsigned block sequences. +//! +//! As with mini-block compression, selectors return concrete codecs. Codecs +//! own their child codecs, framing, and validation. + +use lance_core::{Error, Result}; + +pub(crate) const MAX_DICTIONARY_ITEMS: usize = 4096; +#[cfg(feature = "bitpacking")] +pub(crate) const BITPACK_CHUNK_VALUES: u64 = 1024; + +/// Typed unsigned role expected from a block descriptor. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BlockValueType { + /// Unsigned 8-bit values. + UInt8, + /// Unsigned 16-bit values. + UInt16, + /// Unsigned 32-bit values. + UInt32, + /// Unsigned 64-bit values. + UInt64, +} + +impl BlockValueType { + pub(crate) fn from_bits(bits_per_value: u64) -> Result { + match bits_per_value { + 8 => Ok(Self::UInt8), + 16 => Ok(Self::UInt16), + 32 => Ok(Self::UInt32), + 64 => Ok(Self::UInt64), + _ => Err(Error::invalid_input(format!( + "Block sequence only supports 8, 16, 32, or 64-bit values, got {bits_per_value}" + ))), + } + } + + /// Returns the fixed width of each value. + pub fn bits_per_value(self) -> u64 { + match self { + Self::UInt8 => 8, + Self::UInt16 => 16, + Self::UInt32 => 32, + Self::UInt64 => 64, + } + } + + pub(crate) fn bytes_per_value(self) -> usize { + (self.bits_per_value() / 8) as usize + } + + pub(crate) fn max_value(self) -> u64 { + match self { + Self::UInt8 => u8::MAX as u64, + Self::UInt16 => u16::MAX as u64, + Self::UInt32 => u32::MAX as u64, + Self::UInt64 => u64::MAX, + } + } +} + +mod factory; +pub(crate) mod fixed; + +pub(crate) use factory::{ + create_block_decompressor, encode_scalar, infer_block_value_type, validate_fixed_payload_len, +}; +#[cfg(feature = "bitpacking")] +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 new file mode 100644 index 00000000000..536564a429a --- /dev/null +++ b/rust/lance-encoding/src/compression/block/factory.rs @@ -0,0 +1,670 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Fallible construction of concrete generic block decompressors. + +use bytes::Bytes; + +use super::*; +use crate::{ + buffer::LanceBuffer, + compression::BlockDecompressor, + encodings::physical::{ + block::{CompressionConfig, CompressionScheme}, + constant::ConstantBlockDecompressor, + delta::DeltaDecompressor, + dictionary::DictionaryBlockDecompressor, + general::GenericGeneralBlockDecompressor, + range::RangeDecompressor, + rle::{BlockRleDecompressor, BlockRunCount, MetadataRunLengths}, + value::FixedWidthBlockDecompressor, + }, + format::pb21::{self, CompressiveEncoding, compressive_encoding::Compression}, +}; + +#[cfg(feature = "bitpacking")] +use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking}; + +fn checked_payload_bytes( + value_type: BlockValueType, + num_values: u64, + label: &str, +) -> Result { + let bytes = usize::try_from(num_values) + .ok() + .and_then(|num_values| num_values.checked_mul(value_type.bytes_per_value())) + .ok_or_else(|| Error::invalid_input(format!("{label} payload length overflows usize")))?; + if bytes > isize::MAX as usize { + return Err(Error::invalid_input(format!( + "{label} payload length {bytes} exceeds isize::MAX" + ))); + } + Ok(bytes) +} + +pub fn validate_fixed_payload_len( + payload: &LanceBuffer, + value_type: BlockValueType, + num_values: u64, + label: &str, +) -> Result<()> { + let expected = checked_payload_bytes(value_type, num_values, label)?; + if payload.len() != expected { + return Err(Error::invalid_input(format!( + "{label} payload has {} bytes, expected {expected} for {num_values} {}-bit values", + payload.len(), + value_type.bits_per_value() + ))); + } + Ok(()) +} + +#[cfg(feature = "bitpacking")] +pub fn validate_inline_bitpacking_payload( + payload: &LanceBuffer, + value_type: BlockValueType, + num_values: u64, +) -> Result<()> { + if num_values == 0 || num_values > BITPACK_CHUNK_VALUES { + return Err(Error::invalid_input(format!( + "Inline bitpacking cardinality {num_values} is outside 1..={BITPACK_CHUNK_VALUES}" + ))); + } + let word_bytes = value_type.bytes_per_value(); + if payload.len() < word_bytes { + return Err(Error::invalid_input(format!( + "Inline bitpacking payload has {} bytes, shorter than its {word_bytes}-byte header", + payload.len() + ))); + } + let bit_width = decode_scalar(&payload[..word_bytes], value_type, "Inline bitpacking")?; + if bit_width > value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Inline bitpacking width {bit_width} exceeds {}", + value_type.bits_per_value() + ))); + } + let packed_words = (BITPACK_CHUNK_VALUES * bit_width) / value_type.bits_per_value(); + let expected_words = 1_u64 + .checked_add(packed_words) + .ok_or_else(|| Error::invalid_input("Inline bitpacking payload word count overflows"))?; + let expected_bytes = usize::try_from(expected_words) + .ok() + .and_then(|words| words.checked_mul(word_bytes)) + .ok_or_else(|| Error::invalid_input("Inline bitpacking payload length overflows"))?; + if payload.len() != expected_bytes { + return Err(Error::invalid_input(format!( + "Inline bitpacking payload has {} bytes, expected {expected_bytes}", + payload.len() + ))); + } + Ok(()) +} + +#[cfg(feature = "bitpacking")] +pub fn out_of_line_payload_bytes( + value_type: BlockValueType, + num_values: u64, + compressed_bits_per_value: u64, +) -> Result { + if compressed_bits_per_value >= value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Invalid out-of-line bit width {compressed_bits_per_value} for {}-bit values", + value_type.bits_per_value() + ))); + } + let full_chunks = num_values / BITPACK_CHUNK_VALUES; + let tail_values = num_values % BITPACK_CHUNK_VALUES; + let words_per_chunk = + (BITPACK_CHUNK_VALUES * compressed_bits_per_value).div_ceil(value_type.bits_per_value()); + let mut words = full_chunks + .checked_mul(words_per_chunk) + .ok_or_else(|| Error::invalid_input("Out-of-line bitpacking word count overflows"))?; + if tail_values > 0 { + let tail_bit_savings = value_type.bits_per_value() - compressed_bits_per_value; + let padding_cost = compressed_bits_per_value * (BITPACK_CHUNK_VALUES - tail_values); + let tail_pack_savings = tail_bit_savings * tail_values; + words = words + .checked_add(if padding_cost < tail_pack_savings { + words_per_chunk + } else { + tail_values + }) + .ok_or_else(|| { + Error::invalid_input("Out-of-line bitpacking tail word count overflows") + })?; + } + words + .checked_mul(value_type.bytes_per_value() as u64) + .ok_or_else(|| Error::invalid_input("Out-of-line bitpacking byte length overflows")) +} + +#[cfg(feature = "bitpacking")] +pub fn validate_out_of_line_payload( + payload: &LanceBuffer, + value_type: BlockValueType, + num_values: u64, + compressed_bits_per_value: u64, +) -> Result<()> { + let expected = out_of_line_payload_bytes(value_type, num_values, compressed_bits_per_value)?; + if payload.len() as u64 != expected { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking payload has {} bytes, expected {expected}", + payload.len() + ))); + } + Ok(()) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Position { + Root, + Child, +} + +/// Builds the concrete decoder tree while validating the bounded block grammar. +pub fn create_block_decompressor( + encoding: &CompressiveEncoding, + expected_type: BlockValueType, +) -> Result<(Box, bool)> { + create_inner(encoding, expected_type, Position::Root, true) +} + +fn create_inner( + encoding: &CompressiveEncoding, + expected_type: BlockValueType, + position: Position, + allow_general: bool, +) -> Result<(Box, bool)> { + let compression = encoding + .compression + .as_ref() + .ok_or_else(|| Error::invalid_input("Block encoding is missing its compression variant"))?; + let expected_bits = expected_type.bits_per_value(); + match compression { + Compression::Flat(flat) => { + validate_flat(flat, expected_type, "Flat")?; + Ok(( + Box::new(FixedWidthBlockDecompressor::new(expected_type)), + true, + )) + } + Compression::Constant(constant) => { + let value = constant + .value + .as_deref() + .map(|value| decode_scalar(value, expected_type, "Constant")) + .transpose()?; + Ok(( + Box::new(ConstantBlockDecompressor::new(expected_type, value)), + false, + )) + } + Compression::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, + expected_type, + "Inline bitpacking", + )?; + if bitpacking.values.is_some() { + return Err(Error::invalid_input( + "Inline bitpacking leaf buffer compression is unsupported", + )); + } + #[cfg(feature = "bitpacking")] + { + Ok((Box::new(InlineBitpacking::new(expected_bits)), true)) + } + #[cfg(not(feature = "bitpacking"))] + { + Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )) + } + } + Compression::OutOfLineBitpacking(bitpacking) => { + validate_declared_bits( + bitpacking.uncompressed_bits_per_value, + expected_type, + "Out-of-line bitpacking", + )?; + let values = bitpacking.values.as_deref().ok_or_else(|| { + Error::invalid_input("Out-of-line bitpacking is missing its values encoding") + })?; + let Some(Compression::Flat(flat)) = values.compression.as_ref() else { + return Err(Error::invalid_input( + "Out-of-line bitpacking values must use Flat encoding", + )); + }; + if flat.data.is_some() || flat.bits_per_value >= expected_bits { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking width {} must be between 0 and {}", + flat.bits_per_value, + expected_bits - 1 + ))); + } + #[cfg(feature = "bitpacking")] + { + Ok(( + Box::new(OutOfLineBitpacking::new(flat.bits_per_value, expected_bits)), + true, + )) + } + #[cfg(not(feature = "bitpacking"))] + { + Err(Error::not_supported_source( + "this runtime was not built with bitpacking support".into(), + )) + } + } + Compression::General(general) => { + if position != Position::Root || !allow_general { + return Err(Error::invalid_input( + "General compression is only supported as the single outer block transform", + )); + } + let config = validate_compression_config(general.compression.as_ref(), "General")?; + let child_encoding = general.values.as_deref().ok_or_else(|| { + Error::invalid_input("General compression is missing its child encoding") + })?; + if !matches!( + child_encoding.compression.as_ref(), + Some(Compression::Flat(_)) + ) { + return Err(Error::invalid_input( + "Outer General block compression only supports a Flat child", + )); + } + let (child, child_has_payload) = + create_inner(child_encoding, expected_type, Position::Child, false)?; + if !child_has_payload { + return Err(Error::invalid_input( + "Outer General block compression requires a payload-bearing child", + )); + } + Ok(( + Box::new(GenericGeneralBlockDecompressor::new( + child, + config, + expected_type, + )), + true, + )) + } + Compression::Rle(rle) => { + if position != Position::Root { + return Err(Error::invalid_input( + "RLE is not supported as a block codec child", + )); + } + let values_encoding = rle + .values + .as_deref() + .ok_or_else(|| Error::invalid_input("RLE is missing its values encoding"))?; + let run_lengths_encoding = rle + .run_lengths + .as_deref() + .ok_or_else(|| Error::invalid_input("RLE is missing its run lengths encoding"))?; + let run_length_type = infer_inner(run_lengths_encoding, Position::Child)?; + if !matches!( + run_length_type, + BlockValueType::UInt8 | BlockValueType::UInt16 | BlockValueType::UInt32 + ) { + return Err(Error::invalid_input(format!( + "RLE run lengths must use 8, 16, or 32-bit values, got {}", + run_length_type.bits_per_value() + ))); + } + let (values, values_have_payload) = + create_inner(values_encoding, expected_type, Position::Child, false)?; + let (run_lengths, run_lengths_have_payload) = create_inner( + run_lengths_encoding, + run_length_type, + Position::Child, + false, + )?; + let metadata_run_lengths = match run_lengths_encoding.compression.as_ref() { + Some(Compression::Constant(constant)) => { + let value = decode_scalar( + constant.value.as_deref().ok_or_else(|| { + Error::invalid_input("RLE run lengths Constant is missing its scalar") + })?, + run_length_type, + "RLE run lengths Constant", + )?; + Some(MetadataRunLengths::Constant(value)) + } + 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 { + BlockRunCount::Metadata(metadata) + } else if matches!( + run_lengths_encoding.compression.as_ref(), + Some(Compression::Flat(_)) + ) { + BlockRunCount::RunLengthsPayload + } else if matches!( + values_encoding.compression.as_ref(), + Some(Compression::Flat(_)) + ) { + BlockRunCount::ValuesPayload + } else { + return Err(Error::invalid_input( + "RLE requires metadata run lengths or a Flat child to determine the run count", + )); + }; + Ok(( + Box::new(BlockRleDecompressor::new( + expected_type, + run_length_type, + values, + run_lengths, + values_have_payload, + run_lengths_have_payload, + run_count, + )), + values_have_payload || run_lengths_have_payload, + )) + } + Compression::Dictionary(dictionary) => { + if position != Position::Root { + return Err(Error::invalid_input( + "Dictionary is not supported as a block codec child", + )); + } + if dictionary.num_dictionary_items == 0 + || dictionary.num_dictionary_items as usize > MAX_DICTIONARY_ITEMS + { + return Err(Error::invalid_input(format!( + "Dictionary item count {} is outside 1..={MAX_DICTIONARY_ITEMS}", + dictionary.num_dictionary_items + ))); + } + let indices_encoding = dictionary.indices.as_deref().ok_or_else(|| { + Error::invalid_input("Dictionary is missing its indices encoding") + })?; + let items_encoding = dictionary + .items + .as_deref() + .ok_or_else(|| Error::invalid_input("Dictionary is missing its items encoding"))?; + let (indices, indices_have_payload) = create_inner( + indices_encoding, + BlockValueType::UInt32, + Position::Child, + false, + )?; + let (items, items_have_payload) = + create_inner(items_encoding, expected_type, Position::Child, false)?; + Ok(( + Box::new(DictionaryBlockDecompressor::new( + expected_type, + dictionary.num_dictionary_items, + indices, + items, + indices_have_payload, + items_have_payload, + )), + indices_have_payload || items_have_payload, + )) + } + other => Err(Error::invalid_input(format!( + "Unsupported block sequence encoding: {}", + compression_name(other) + ))), + } +} + +fn validate_flat(flat: &pb21::Flat, expected_type: BlockValueType, label: &str) -> Result<()> { + validate_declared_bits(flat.bits_per_value, expected_type, label)?; + if flat.data.is_some() { + return Err(Error::invalid_input(format!( + "{label} leaf buffer compression is unsupported" + ))); + } + Ok(()) +} + +fn validate_declared_bits(actual: u64, expected_type: BlockValueType, label: &str) -> Result<()> { + if actual != expected_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "{label} declares {actual}-bit values, expected {}", + expected_type.bits_per_value() + ))); + } + Ok(()) +} + +fn decode_scalar(bytes: &[u8], value_type: BlockValueType, label: &str) -> Result { + if bytes.len() != value_type.bytes_per_value() { + return Err(Error::invalid_input(format!( + "{label} scalar has {} bytes, expected {}", + bytes.len(), + value_type.bytes_per_value() + ))); + } + Ok(match value_type { + BlockValueType::UInt8 => u64::from(bytes[0]), + BlockValueType::UInt16 => u64::from(u16::from_le_bytes([bytes[0], bytes[1]])), + BlockValueType::UInt32 => u64::from(u32::from_le_bytes( + bytes.try_into().expect("scalar length was checked"), + )), + BlockValueType::UInt64 => { + u64::from_le_bytes(bytes.try_into().expect("scalar length was checked")) + } + }) +} + +pub fn encode_scalar(value: u64, value_type: BlockValueType) -> Result { + if value > value_type.max_value() { + return Err(Error::invalid_input(format!( + "Scalar value {value} exceeds the {}-bit value range", + value_type.bits_per_value() + ))); + } + Ok(match value_type { + BlockValueType::UInt8 => vec![value as u8], + BlockValueType::UInt16 => (value as u16).to_le_bytes().to_vec(), + BlockValueType::UInt32 => (value as u32).to_le_bytes().to_vec(), + BlockValueType::UInt64 => value.to_le_bytes().to_vec(), + } + .into()) +} + +fn validate_compression_config( + compression: Option<&pb21::BufferCompression>, + label: &str, +) -> Result { + let compression = compression + .ok_or_else(|| Error::invalid_input(format!("{label} is missing compression config")))?; + let scheme = pb21::CompressionScheme::try_from(compression.scheme).map_err(|_| { + Error::invalid_input(format!( + "{label} has unknown compression scheme {}", + compression.scheme + )) + })?; + let scheme = CompressionScheme::try_from(scheme)?; + Ok(CompressionConfig::new(scheme, compression.level)) +} + +/// 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 + .as_ref() + .ok_or_else(|| Error::invalid_input("Block encoding is missing its compression variant"))?; + match compression { + Compression::Flat(flat) => BlockValueType::from_bits(flat.bits_per_value), + Compression::Constant(constant) => { + let value = constant.value.as_ref().ok_or_else(|| { + Error::invalid_input( + "Cannot infer an empty Constant block type without its typed container role", + ) + })?; + BlockValueType::from_bits((value.len() * 8) as u64) + } + Compression::InlineBitpacking(bitpacking) => { + BlockValueType::from_bits(bitpacking.uncompressed_bits_per_value) + } + Compression::OutOfLineBitpacking(bitpacking) => { + BlockValueType::from_bits(bitpacking.uncompressed_bits_per_value) + } + Compression::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 + .as_deref() + .ok_or_else(|| Error::invalid_input("General is missing its child encoding"))?, + Position::Child, + ), + Compression::Rle(rle) if position == Position::Root => infer_inner( + rle.values + .as_deref() + .ok_or_else(|| Error::invalid_input("RLE is missing its values encoding"))?, + Position::Child, + ), + Compression::Dictionary(dictionary) if position == Position::Root => infer_inner( + dictionary + .items + .as_deref() + .ok_or_else(|| Error::invalid_input("Dictionary is missing its items encoding"))?, + Position::Child, + ), + other => Err(Error::invalid_input(format!( + "Cannot infer bounded block value type from {} at {position:?}", + compression_name(other) + ))), + } +} + +fn compression_name(compression: &Compression) -> &'static str { + match compression { + Compression::Flat(_) => "flat", + Compression::Variable(_) => "variable", + Compression::Constant(_) => "constant", + Compression::OutOfLineBitpacking(_) => "out-of-line bitpacking", + Compression::InlineBitpacking(_) => "inline bitpacking", + Compression::Fsst(_) => "fsst", + Compression::Dictionary(_) => "dictionary", + Compression::Rle(_) => "rle", + Compression::ByteStreamSplit(_) => "byte stream split", + Compression::General(_) => "general", + Compression::FixedSizeList(_) => "fixed-size list", + Compression::PackedStruct(_) => "packed struct", + Compression::VariablePackedStruct(_) => "variable packed struct", + Compression::Range(_) => "range", + Compression::Delta(_) => "delta", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::format::ProtobufUtils21; + + #[test] + fn rejects_nested_transform_trees() { + let nested = ProtobufUtils21::delta( + 64, + 0, + ProtobufUtils21::delta(64, 0, ProtobufUtils21::flat(64, None)), + ); + assert!( + create_block_decompressor(&nested, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("child") + ); + } + + #[test] + fn validates_range_cardinality_when_decoding() { + let encoding = 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/fixed.rs b/rust/lance-encoding/src/compression/block/fixed.rs new file mode 100644 index 00000000000..409bd5a0c3d --- /dev/null +++ b/rust/lance-encoding/src/compression/block/fixed.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use super::*; +use crate::{ + buffer::LanceBuffer, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, +}; + +pub fn fixed_block(value_type: BlockValueType, num_values: u64, data: LanceBuffer) -> DataBlock { + DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: value_type.bits_per_value(), + data, + num_values, + block_info: BlockInfo::default(), + }) +} + +pub fn fixed_from_u64_values( + values: &[u64], + value_type: BlockValueType, + label: &str, +) -> Result { + let data = match value_type { + BlockValueType::UInt8 => LanceBuffer::reinterpret_vec( + values + .iter() + .map(|value| { + u8::try_from(*value).map_err(|_| { + Error::invalid_input(format!("{label} value {value} exceeds u8::MAX")) + }) + }) + .collect::>>()?, + ), + BlockValueType::UInt16 => LanceBuffer::reinterpret_vec( + values + .iter() + .map(|value| { + u16::try_from(*value).map_err(|_| { + Error::invalid_input(format!("{label} value {value} exceeds u16::MAX")) + }) + }) + .collect::>>()?, + ), + BlockValueType::UInt32 => LanceBuffer::reinterpret_vec( + values + .iter() + .map(|value| { + u32::try_from(*value).map_err(|_| { + Error::invalid_input(format!("{label} value {value} exceeds u32::MAX")) + }) + }) + .collect::>>()?, + ), + BlockValueType::UInt64 => LanceBuffer::reinterpret_vec(values.to_vec()), + }; + Ok(FixedWidthDataBlock { + bits_per_value: value_type.bits_per_value(), + data, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + }) +} + +pub fn visit_unsigned_values( + block: &FixedWidthDataBlock, + value_type: BlockValueType, + mut visit: impl FnMut(u64) -> Result<()>, +) -> Result<()> { + validate_fixed_payload_len(&block.data, value_type, block.num_values, "Block input")?; + match value_type { + BlockValueType::UInt8 => { + for value in block.data.iter().copied() { + visit(u64::from(value))?; + } + } + BlockValueType::UInt16 => { + for value in block.data.borrow_to_typed_view::().iter().copied() { + visit(u64::from(value))?; + } + } + BlockValueType::UInt32 => { + for value in block.data.borrow_to_typed_view::().iter().copied() { + visit(u64::from(value))?; + } + } + BlockValueType::UInt64 => { + for value in block.data.borrow_to_typed_view::().iter().copied() { + visit(value)?; + } + } + } + Ok(()) +} + +pub fn read_unsigned_values( + block: &FixedWidthDataBlock, + value_type: BlockValueType, +) -> Result> { + validate_fixed_payload_len( + &block.data, + value_type, + block.num_values, + "Fixed-width block", + )?; + Ok(match value_type { + BlockValueType::UInt8 => block.data.iter().map(|value| u64::from(*value)).collect(), + BlockValueType::UInt16 => block + .data + .borrow_to_typed_slice::() + .iter() + .map(|value| u64::from(*value)) + .collect(), + BlockValueType::UInt32 => block + .data + .borrow_to_typed_slice::() + .iter() + .map(|value| u64::from(*value)) + .collect(), + BlockValueType::UInt64 => block.data.borrow_to_typed_slice::().to_vec(), + }) +} diff --git a/rust/lance-encoding/src/compression/block/tests.rs b/rust/lance-encoding/src/compression/block/tests.rs new file mode 100644 index 00000000000..230f9660ee8 --- /dev/null +++ b/rust/lance-encoding/src/compression/block/tests.rs @@ -0,0 +1,413 @@ +use super::*; + +use std::sync::Arc; + +use crate::{ + buffer::LanceBuffer, + compression::BlockCompressor, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + encodings::physical::{ + constant::ConstantBlockCompressor, dictionary::DictionaryBlockCompressor, + range::RangeEncoder, rle::BlockRleCompressor, value::FixedWidthBlockCompressor, + }, + format::{ProtobufUtils21, pb21::CompressiveEncoding}, +}; + +fn fixed_u64(values: &[u64]) -> FixedWidthDataBlock { + FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(values.to_vec()), + bits_per_value: 64, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + } +} + +fn decoded_u64(block: DataBlock) -> Vec { + let DataBlock::FixedWidth(block) = block else { + panic!("expected fixed-width output"); + }; + block.data.borrow_to_typed_slice::().to_vec() +} + +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() +} + +fn round_trip_u64( + values: &[u64], + compressor: Box, + encoding: CompressiveEncoding, +) { + let payload = compressor + .compress(DataBlock::FixedWidth(fixed_u64(values))) + .unwrap(); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert_eq!(payload.is_some(), has_payload); + assert_eq!( + decoded_u64(decoder.decompress(payload, values.len() as u64).unwrap()), + values + ); +} + +#[test] +fn metadata_compressors_validate_and_decode() { + round_trip_u64( + &[], + Box::new(ConstantBlockCompressor::new(BlockValueType::UInt64, None)), + ProtobufUtils21::constant(None), + ); + round_trip_u64( + &[7_u64; 32], + Box::new(ConstantBlockCompressor::new( + BlockValueType::UInt64, + Some(7), + )), + ProtobufUtils21::constant(Some(7_u64.to_le_bytes().to_vec().into())), + ); + 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] +fn payload_presence_distinguishes_metadata_from_empty_payload() { + let flat = Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)); + let payload = flat + .compress(DataBlock::FixedWidth(fixed_u64(&[]))) + .unwrap(); + assert!(payload.as_ref().is_some_and(|payload| payload.is_empty())); + + let (flat_decoder, flat_has_payload) = + create_block_decompressor(&ProtobufUtils21::flat(64, None), BlockValueType::UInt64) + .unwrap(); + assert!(flat_has_payload); + assert!(flat_decoder.decompress(None, 0).is_err()); + assert!( + decoded_u64( + flat_decoder + .decompress(Some(LanceBuffer::empty()), 0) + .unwrap() + ) + .is_empty() + ); + + let (constant_decoder, constant_has_payload) = create_block_decompressor( + &ProtobufUtils21::constant(Some(7_u64.to_le_bytes().to_vec().into())), + BlockValueType::UInt64, + ) + .unwrap(); + assert!(!constant_has_payload); + assert!( + constant_decoder + .decompress(Some(LanceBuffer::empty()), 1) + .is_err() + ); +} + +#[test] +fn rle_compressor_owns_and_reuses_children() { + let mut values = Vec::new(); + for value in [11_u64, 91, 37, 123] { + values.extend(std::iter::repeat_n(value, 256)); + } + let compressor = Box::new(BlockRleCompressor::new( + BlockValueType::UInt64, + BlockValueType::UInt16, + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)), + Box::new(ConstantBlockCompressor::new( + BlockValueType::UInt16, + Some(256), + )), + )); + let encoding = ProtobufUtils21::rle( + ProtobufUtils21::flat(64, None), + ProtobufUtils21::constant(Some(256_u16.to_le_bytes().to_vec().into())), + ); + round_trip_u64(&values, compressor, encoding); +} + +#[test] +fn dictionary_materializes_once_when_built() { + use crate::encodings::physical::dictionary::BLOCK_MATERIALIZATION_COUNT; + + let values = (0..4096_u64) + .map(|index| [u64::MAX - 9, 17, 1_u64 << 50, 991][index as usize % 4]) + .collect::>(); + let dictionary_items = Arc::from([17_u64, 991, 1_u64 << 50, u64::MAX - 9].as_slice()); + let compressor = Box::new(DictionaryBlockCompressor::new( + BlockValueType::UInt64, + dictionary_items, + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt32)), + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)), + )); + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::flat(64, None), + 4, + ); + BLOCK_MATERIALIZATION_COUNT.with(|count| count.set(0)); + BLOCK_MATERIALIZATION_COUNT.with(|count| assert_eq!(count.get(), 0)); + round_trip_u64(&values, compressor, encoding); + BLOCK_MATERIALIZATION_COUNT.with(|count| assert_eq!(count.get(), 1)); +} + +#[cfg(feature = "bitpacking")] +#[test] +fn out_of_line_bitpacking_round_trip() { + use crate::encodings::physical::bitpacking::OutOfLineBitpacking; + + let values = (0..4096_u64) + .map(|index| (index * 37) % 1024) + .collect::>(); + round_trip_u64( + &values, + Box::new(OutOfLineBitpacking::new(10, 64)), + ProtobufUtils21::out_of_line_bitpacking(64, ProtobufUtils21::flat(10, None)), + ); +} + +#[cfg(any(feature = "lz4", feature = "zstd"))] +#[test] +fn general_compressor_owns_a_flat_child() { + use crate::encodings::physical::{ + block::{CompressionConfig, CompressionScheme}, + general::GeneralBlockCompressor, + }; + + let scheme = if cfg!(feature = "lz4") { + CompressionScheme::Lz4 + } else { + CompressionScheme::Zstd + }; + let config = CompressionConfig::new(scheme, None); + let compressor = Box::new(GeneralBlockCompressor::new( + Box::new(FixedWidthBlockCompressor::new(BlockValueType::UInt64)), + config, + )); + let encoding = ProtobufUtils21::wrapped(config, ProtobufUtils21::flat(64, None)).unwrap(); + let values = (0..16_384_u64).map(|value| value % 17).collect::>(); + round_trip_u64(&values, compressor, encoding); +} + +#[test] +fn factory_rejects_unbounded_or_mistyped_trees() { + let nested_delta = ProtobufUtils21::delta( + 64, + 0, + ProtobufUtils21::delta(64, 0, ProtobufUtils21::flat(64, None)), + ); + assert!( + create_block_decompressor(&nested_delta, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("child") + ); + + let wrong_width = ProtobufUtils21::range(32, 0, 1); + assert!( + create_block_decompressor(&wrong_width, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("expected 64") + ); + + let oversized = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::flat(64, None), + (MAX_DICTIONARY_ITEMS + 1) as u32, + ); + assert!( + create_block_decompressor(&oversized, BlockValueType::UInt64) + .unwrap_err() + .to_string() + .contains("outside") + ); +} + +#[test] +fn constant_cardinality_contract_is_checked_at_decode() { + let (empty, has_payload) = + create_block_decompressor(&ProtobufUtils21::constant(None), BlockValueType::UInt64) + .unwrap(); + assert!(!has_payload); + assert!(decoded_u64(empty.decompress(None, 0).unwrap()).is_empty()); + assert!(empty.decompress(None, 1).is_err()); + + let (present, has_payload) = create_block_decompressor( + &ProtobufUtils21::constant(Some(7_u64.to_le_bytes().to_vec().into())), + BlockValueType::UInt64, + ) + .unwrap(); + assert!(!has_payload); + assert!(present.decompress(None, 0).is_err()); + assert_eq!( + decoded_u64(present.decompress(None, 3).unwrap()), + vec![7, 7, 7] + ); +} + +#[test] +fn 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( + 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( + ProtobufUtils21::flat(64, None), + ProtobufUtils21::constant(Some(vec![2_u8].into())), + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(has_payload); + assert!( + decoder + .decompress(Some(LanceBuffer::from(vec![0; 7])), 4) + .is_err() + ); + + let mut payload = 100_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&[0; 8]); + assert!( + decoder + .decompress(Some(LanceBuffer::from(payload)), 4) + .is_err() + ); + + let mut payload = 16_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&[0; 16]); + payload.push(1); + let error = decoder + .decompress(Some(LanceBuffer::from(payload)), 4) + .unwrap_err(); + assert!(error.to_string().contains("Metadata-only RLE run-length")); +} + +#[test] +fn metadata_dictionary_checks_index_bounds() { + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::range(32, 0, 1), + ProtobufUtils21::range(64, 10, 1), + 2, + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(!has_payload); + let error = decoder.decompress(None, 3).unwrap_err(); + assert!(error.to_string().contains("out of bounds")); +} + +#[test] +fn dictionary_framing_rejects_inconsistent_lengths() { + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::flat(64, None), + 2, + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(has_payload); + let mut payload = 4_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&16_u64.to_le_bytes()); + payload.extend_from_slice(&[0; 8]); + assert!( + decoder + .decompress(Some(LanceBuffer::from(payload)), 1) + .is_err() + ); +} + +#[test] +fn dictionary_framing_rejects_payload_for_metadata_child() { + let encoding = ProtobufUtils21::dictionary( + ProtobufUtils21::flat(32, None), + ProtobufUtils21::range(64, 10, 1), + 2, + ); + let (decoder, has_payload) = + create_block_decompressor(&encoding, BlockValueType::UInt64).unwrap(); + assert!(has_payload); + let mut payload = 4_u64.to_le_bytes().to_vec(); + payload.extend_from_slice(&1_u64.to_le_bytes()); + payload.extend_from_slice(&0_u32.to_le_bytes()); + payload.push(0); + let error = decoder + .decompress(Some(LanceBuffer::from(payload)), 1) + .unwrap_err(); + assert!(error.to_string().contains("Metadata-only Dictionary items")); +} + +#[test] +fn decoder_allocation_overflow_is_fallible() { + let (decoder, has_payload) = create_block_decompressor( + &ProtobufUtils21::constant(Some(1_u64.to_le_bytes().to_vec().into())), + BlockValueType::UInt64, + ) + .unwrap(); + assert!(!has_payload); + assert!(decoder.decompress(None, u64::MAX).is_err()); +} + +#[cfg(feature = "bitpacking")] +#[test] +fn frozen_out_of_line_compressor_rejects_wider_reuse() { + use crate::encodings::physical::bitpacking::OutOfLineBitpacking; + + let compressor = OutOfLineBitpacking::new(3, 64); + let error = compressor + .compress(DataBlock::FixedWidth(fixed_u64(&[1, 2, 8]))) + .unwrap_err(); + assert!(error.to_string().contains("requires 4 bits")); +} diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 1b16ab73b80..2853932b352 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -45,7 +45,7 @@ use crate::utils::bytepack::ByteUnpacker; use crate::{ compression::{ BlockDecompressor, CompressionStrategy, DecompressionStrategy, MiniBlockDecompressor, - create_rle_decompressor, + compress_required_block, create_rle_decompressor, }, data::{AllNullDataBlock, DataBlock, VariableWidthBlock}, utils::bytepack::BytepackedIntegerEncoder, @@ -55,7 +55,7 @@ use crate::{ encodings::logical::primitive::fullzip::PerValueDataBlock, }; use crate::{ - encodings::logical::primitive::miniblock::MiniBlockCompressed, + encodings::logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, statistics::{ComputeStat, GetStat, Stat}, }; use crate::{ @@ -179,7 +179,7 @@ impl DecodeMiniBlockTask { levels: LanceBuffer, num_levels: u16, ) -> Result> { - let rep = rep_decompressor.decompress(levels, num_levels as u64)?; + let rep = rep_decompressor.decompress(Some(levels), num_levels as u64)?; let rep = rep.as_fixed_width().unwrap(); debug_assert_eq!(rep.num_values, num_levels as u64); debug_assert_eq!(rep.bits_per_value, 16); @@ -1556,7 +1556,7 @@ impl StructuralPageScheduler for ComplexAllNullScheduler { } LevelCodec::Block(decompressor) => { let frame = LanceBuffer::from_bytes(compressed_bytes, 1); - let decompressed = decompressor.decompress(frame, num_values)?; + let decompressed = decompressor.decompress(Some(frame), num_values)?; dense_levels_from_block(decompressed, num_values, level_type) } } @@ -2580,7 +2580,10 @@ impl StructuralPageScheduler for MiniBlockScheduler { let dictionary = if let Some(ref mut dictionary) = self.dictionary { let dictionary_data = dictionary_bytes.unwrap(); Some(Arc::new(dictionary.dictionary_decompressor.decompress( - LanceBuffer::from_bytes(dictionary_data, dictionary.dictionary_data_alignment), + Some(LanceBuffer::from_bytes( + dictionary_data, + dictionary.dictionary_data_alignment, + )), dictionary.num_dictionary_items, )?)) } else { @@ -4827,7 +4830,6 @@ impl PrimitiveStructuralEncoder { let levels_block = DataBlock::FixedWidth(fixed_width_block); let levels_field = Field::new_arrow("", DataType::UInt16, false)?; - // Pick a block compressor let (compressor, compressor_desc) = compression_strategy.create_block_compressor(&levels_field, &levels_block)?; // Compress blocks of levels (sized according to the chunks) @@ -4899,7 +4901,11 @@ impl PrimitiveStructuralEncoder { }; chunk_fixed_width.compute_stat(); let chunk_levels_block = DataBlock::FixedWidth(chunk_fixed_width); - let compressed_levels = compressor.compress(chunk_levels_block)?; + let compressed_levels = compressor.compress(chunk_levels_block)?.ok_or_else(|| { + Error::internal( + "Rep/def block compressor selected a metadata-only codec".to_string(), + ) + })?; let num_levels = u16::try_from(num_chunk_levels).map_err(|_| { Error::invalid_input_source( format!( @@ -4960,9 +4966,8 @@ impl PrimitiveStructuralEncoder { let levels_block = DataBlock::FixedWidth(fixed_width_block); let levels_field = Field::new_arrow("", DataType::UInt16, false)?; - let (compressor, encoding) = - compression_strategy.create_block_compressor(&levels_field, &levels_block)?; - let compressed_buffer = compressor.compress(levels_block)?; + let (compressed_buffer, encoding) = + compress_required_block(compression_strategy, &levels_field, levels_block)?; Ok((compressed_buffer, encoding)) } @@ -5249,7 +5254,7 @@ impl PrimitiveStructuralEncoder { fn build_dict_values_compressor_field(field: &Field) -> Result { // This is an internal synthetic field used only to feed metadata into - // `create_block_compressor` for dictionary values. The concrete type/name here + // block planning for dictionary values. The concrete type/name here // are not semantically meaningful; we rely on explicit metadata below to control // general compression selection for dictionary values. let mut dict_values_field = Field::new_arrow("", DataType::UInt16, false)?; @@ -5282,7 +5287,11 @@ impl PrimitiveStructuralEncoder { let num_items = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; - let (compressed_data, value_encoding) = compressor.compress(data)?; + let common_chunk_buffers = + u64::from(repdef.rep_slicer().is_some()) + u64::from(repdef.def_slicer().is_some()); + let compression_context = + MiniBlockCompressionContext::new(common_chunk_buffers, support_large_chunk, true); + let (compressed_data, value_encoding) = compressor.compress(data, compression_context)?; let max_rep = repdef.def_meaning.iter().filter(|l| l.is_list()).count() as u16; @@ -5342,9 +5351,8 @@ impl PrimitiveStructuralEncoder { let num_dictionary_items = dictionary_data.num_values(); let dict_values_field = Self::build_dict_values_compressor_field(field)?; - let (compressor, dictionary_encoding) = compression_strategy - .create_block_compressor(&dict_values_field, &dictionary_data)?; - let dictionary_buffer = compressor.compress(dictionary_data)?; + let (dictionary_buffer, dictionary_encoding) = + compress_required_block(compression_strategy, &dict_values_field, dictionary_data)?; data.push(dictionary_buffer); if let Some(rep_index) = rep_index { @@ -6593,16 +6601,18 @@ mod tests { ChunkInstructions, DataBlock, DecodeMiniBlockTask, FixedPerValueDecompressor, FixedWidthDataBlock, FullZipCacheableState, FullZipDecodeDetails, FullZipReadSource, FullZipRepIndexDetails, FullZipScheduler, LazyLevels, LevelCodec, LevelCursor, LevelPlan, - MiniBlockChunk, MiniBlockChunkIndex, MiniBlockCompressed, PerValueDecompressor, - PreambleAction, RunEndsBuilder, RunPosition, RunStorage, StructuralPageScheduler, - VariableFullZipDecoder, dense_levels_from_block, validate_complex_all_null_levels, + MiniBlockChunk, MiniBlockChunkIndex, MiniBlockCompressed, MiniBlockScheduler, + PerValueDecompressor, PreambleAction, RunEndsBuilder, RunPosition, RunStorage, + StructuralPageScheduler, VariableFullZipDecoder, dense_levels_from_block, + validate_complex_all_null_levels, }; + use crate::EncodingsIo; use crate::buffer::LanceBuffer; use crate::compression::{BlockCompressor, DefaultDecompressionStrategy}; use crate::constants::{ - COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_VALUES_COMPRESSION_LEVEL_META_KEY, - DICT_VALUES_COMPRESSION_META_KEY, STRUCTURAL_ENCODING_META_KEY, - STRUCTURAL_ENCODING_MINIBLOCK, + COMPRESSION_LEVEL_META_KEY, COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY, + DICT_VALUES_COMPRESSION_LEVEL_META_KEY, DICT_VALUES_COMPRESSION_META_KEY, + STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, }; use crate::data::BlockInfo; use crate::decoder::{PageEncoding, StructuralFieldDecoder}; @@ -6619,8 +6629,14 @@ mod tests { use arrow_array::{ArrayRef, Int8Array, StringArray}; use arrow_buffer::ScalarBuffer; use arrow_schema::{DataType, Field as ArrowField}; + use futures::{FutureExt, future::BoxFuture}; + use prost::Message; use std::collections::HashMap; - use std::{collections::VecDeque, sync::Arc}; + use std::{ + collections::VecDeque, + ops::Range, + sync::{Arc, Mutex}, + }; #[test] fn test_is_narrow() { @@ -8672,6 +8688,268 @@ mod tests { pages.into_iter().next().unwrap() } + fn variable_offset_test_field() -> arrow_schema::Field { + arrow_schema::Field::new("c", DataType::Utf8, false).with_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()), + ])) + } + + fn variable_offset_test_array(lengths: &[usize], num_rows: usize) -> ArrayRef { + Arc::new(StringArray::from_iter_values((0..num_rows).map(|index| { + let length = lengths[index % lengths.len()]; + assert!(length >= 4); + format!("{index:04x}{}", "x".repeat(length - 4)) + }))) + } + + fn miniblock_layout(page: &crate::encoder::EncodedPage) -> &pb21::MiniBlockLayout { + let PageEncoding::Structural(layout) = &page.description else { + panic!("Expected structural page encoding"); + }; + let pb21::page_layout::Layout::MiniBlockLayout(layout) = layout.layout.as_ref().unwrap() + else { + panic!("Expected mini-block page layout"); + }; + layout + } + + fn variable_offset_compression( + page: &crate::encoder::EncodedPage, + ) -> &pb21::compressive_encoding::Compression { + let value_encoding = miniblock_layout(page).value_compression.as_ref().unwrap(); + let Compression::Variable(variable) = value_encoding.compression.as_ref().unwrap() else { + panic!("Expected Variable value compression"); + }; + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()) + .unwrap() + } + + fn variable_value_wire_bytes(page: &crate::encoder::EncodedPage) -> usize { + let descriptor_bytes = miniblock_layout(page) + .value_compression + .as_ref() + .unwrap() + .encoded_len(); + descriptor_bytes + page.data.iter().map(LanceBuffer::len).sum::() + } + + #[derive(Debug)] + struct RecordingScheduler { + data: Bytes, + requests: Mutex>>>, + } + + impl RecordingScheduler { + fn new(data: Bytes) -> Self { + Self { + data, + requests: Mutex::new(Vec::new()), + } + } + + fn take_requests(&self) -> Vec>> { + std::mem::take(&mut *self.requests.lock().unwrap()) + } + } + + impl EncodingsIo for RecordingScheduler { + fn submit_request( + &self, + ranges: Vec>, + _priority: u64, + ) -> BoxFuture<'static, lance_core::Result>> { + self.requests.lock().unwrap().push(ranges.clone()); + let data = ranges + .into_iter() + .map(|range| self.data.slice(range.start as usize..range.end as usize)) + .collect(); + std::future::ready(Ok(data)).boxed() + } + } + + async fn miniblock_take_request_shape(page: &crate::encoder::EncodedPage) -> [usize; 6] { + let mut position = 0_u64; + let buffer_offsets_and_sizes = page + .data + .iter() + .map(|buffer| { + let size = buffer.len() as u64; + let descriptor = (position, size); + position += size; + descriptor + }) + .collect::>(); + let mut bytes = Vec::with_capacity(position as usize); + for buffer in &page.data { + bytes.extend_from_slice(buffer); + } + + let recorder = Arc::new(RecordingScheduler::new(Bytes::from(bytes))); + let io: Arc = recorder.clone(); + let decompression = DefaultDecompressionStrategy::default(); + let mut cold = MiniBlockScheduler::try_new( + &buffer_offsets_and_sizes, + 0, + page.num_rows, + miniblock_layout(page), + &decompression, + ) + .unwrap(); + let cached = cold.initialize(&io).await.unwrap(); + let initialize_requests = recorder.take_requests(); + + let cold_tasks = cold.schedule_ranges(&[123..124], &io).unwrap(); + let cold_requests = recorder.take_requests(); + for task in cold_tasks { + task.decoder_fut.await.unwrap(); + } + + let mut warm = MiniBlockScheduler::try_new( + &buffer_offsets_and_sizes, + 0, + page.num_rows, + miniblock_layout(page), + &decompression, + ) + .unwrap(); + warm.load(&cached); + let warm_tasks = warm.schedule_ranges(&[123..124], &io).unwrap(); + let warm_requests = recorder.take_requests(); + for task in warm_tasks { + task.decoder_fut.await.unwrap(); + } + + [ + initialize_requests.len(), + initialize_requests.iter().map(Vec::len).sum(), + cold_requests.len(), + cold_requests.iter().map(Vec::len).sum(), + warm_requests.len(), + warm_requests.iter().map(Vec::len).sum(), + ] + } + + #[tokio::test] + async fn test_v2_3_variable_offsets_use_complete_serialized_cost() { + let cases = [ + (&[16_usize, 22, 17, 20][..], "legacy"), + (&[4_usize, 10, 5, 8][..], "delta"), + (&[8_usize][..], "range"), + ]; + for (lengths, expected) in cases { + let array = variable_offset_test_array(lengths, 10_000); + let legacy = encode_first_page( + variable_offset_test_field(), + array.clone(), + LanceFileVersion::V2_2, + ) + .await; + let selected = + encode_first_page(variable_offset_test_field(), array, LanceFileVersion::V2_3) + .await; + + match expected { + "legacy" => { + assert!(matches!( + variable_offset_compression(&selected), + Compression::Flat(_) + )); + assert_eq!( + variable_value_wire_bytes(&selected), + variable_value_wire_bytes(&legacy) + ); + } + "delta" => assert!(matches!( + variable_offset_compression(&selected), + Compression::Delta(_) + )), + "range" => assert!(matches!( + variable_offset_compression(&selected), + Compression::Range(_) + )), + _ => unreachable!(), + } + if expected != "legacy" { + assert!( + variable_value_wire_bytes(&selected) < variable_value_wire_bytes(&legacy), + "{expected} generic container must be strictly smaller than legacy" + ); + } + } + } + + #[tokio::test] + async fn test_v2_3_variable_offsets_use_delta_range_for_increasing_lengths() { + let array = Arc::new(StringArray::from_iter_values( + (0..262_144).map(|index| "x".repeat(4 + index % 64)), + )); + let legacy = encode_first_page( + variable_offset_test_field(), + array.clone(), + LanceFileVersion::V2_2, + ) + .await; + let selected = + encode_first_page(variable_offset_test_field(), array, LanceFileVersion::V2_3).await; + + let Compression::Delta(delta) = variable_offset_compression(&selected) else { + panic!( + "expected Delta offsets, got {:?}", + variable_offset_compression(&selected) + ); + }; + assert!(matches!( + delta + .deltas + .as_deref() + .and_then(|deltas| deltas.compression.as_ref()), + Some(Compression::Range(_)) + )); + assert!( + variable_value_wire_bytes(&selected) < variable_value_wire_bytes(&legacy), + "Delta(Range) generic container must be strictly smaller than legacy" + ); + } + + #[tokio::test] + async fn test_v2_3_generic_offsets_do_not_add_take_requests() { + let array = variable_offset_test_array(&[4, 10, 5, 8], 10_000); + let legacy = encode_first_page( + variable_offset_test_field(), + array.clone(), + LanceFileVersion::V2_2, + ) + .await; + let generic = + encode_first_page(variable_offset_test_field(), array, LanceFileVersion::V2_3).await; + + assert!(matches!( + variable_offset_compression(&legacy), + Compression::Flat(_) + )); + assert!(matches!( + variable_offset_compression(&generic), + Compression::Delta(_) + )); + assert_eq!(miniblock_layout(&legacy).num_buffers, 1); + assert_eq!(miniblock_layout(&generic).num_buffers, 2); + assert_eq!(legacy.data.len(), 2); + assert_eq!(generic.data.len(), 2); + + let legacy_requests = miniblock_take_request_shape(&legacy).await; + let generic_requests = miniblock_take_request_shape(&generic).await; + assert_eq!(legacy_requests, [1, 1, 1, 1, 1, 1]); + assert_eq!(generic_requests, legacy_requests); + } + #[tokio::test] async fn test_constant_layout_out_of_line_fixed_size_binary_v2_2() { use crate::format::pb21::page_layout::Layout; @@ -8910,7 +9188,7 @@ mod tests { .create_block_decompressor(&encoding) .unwrap(); let decompressed = decompressor - .decompress(compressed_buf, values.len() as u64) + .decompress(Some(compressed_buf), values.len() as u64) .unwrap(); let decompressed_fixed_width = decompressed.as_fixed_width().unwrap(); assert_eq!(decompressed_fixed_width.num_values, values.len() as u64); @@ -8998,6 +9276,7 @@ mod tests { }); BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block) .unwrap() + .unwrap() } fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns { diff --git a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs index edfba526670..1d25d1fdfb0 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/miniblock.rs @@ -54,6 +54,42 @@ pub struct MiniBlockCompressed { pub num_values: u64, } +/// Per-page framing details that can affect a mini-block compressor's choice. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct MiniBlockCompressionContext { + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, +} + +impl MiniBlockCompressionContext { + /// Creates the framing context supplied by the owning mini-block page. + pub fn new( + common_chunk_buffers: u64, + support_large_chunk: bool, + allow_generic_offsets: bool, + ) -> Self { + Self { + common_chunk_buffers, + support_large_chunk, + allow_generic_offsets, + } + } + + /// 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 /// /// Mini-block chunks are designed to be small (just a few disk sectors) @@ -113,7 +149,11 @@ pub trait MiniBlockCompressor: std::fmt::Debug + Send + Sync { /// /// This method also returns a description of the encoding applied that will be /// used at decode time to read the data. - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; + fn compress( + &self, + page: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)>; } #[cfg(test)] @@ -129,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/logical/primitive/sparse.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs index 68becc0014f..d38c32d7ecb 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse.rs @@ -1069,7 +1069,6 @@ enum SparsePositionSetDecoder { }, Explicit { decompressor: Arc, - encoding: CompressiveEncoding, count: u64, domain_len: u64, }, @@ -1084,7 +1083,6 @@ enum SparseCountSetDecoder { }, Explicit { decompressor: Arc, - encoding: CompressiveEncoding, count: u64, }, } @@ -1297,7 +1295,7 @@ impl SparseStructuralScheduler { let layer_decompressors = layout .structural_layers .iter() - .map(|layer| Self::layer_decompressors(layer, decompressors)) + .map(Self::layer_decompressors) .collect::>>()?; Ok(Self { @@ -1388,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(_) => {} } } @@ -1537,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(), )), } @@ -2095,35 +2097,30 @@ impl SparseStructuralScheduler { fn create_position_decompressor( compression: &CompressiveEncoding, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result> { let compression = Self::validate_compression(compression, label)?; - Self::validate_block_encoding(compression, label)?; - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - decompressors.create_block_decompressor(compression) - })) - .map_err(|_| { - Error::invalid_input_source( - format!( - "Sparse structural {label} descriptor caused decompressor construction to panic" - ) - .into(), - ) - })? - .map(Arc::from) + let (decompressor, has_payload) = crate::compression::block::create_block_decompressor( + compression, + crate::compression::BlockValueType::UInt64, + ) .map_err(|error| { Error::invalid_input_source( - format!("Sparse structural {label} decompressor construction failed: {error}") - .into(), + format!("Sparse structural {label} block validation failed: {error}").into(), ) - }) + })?; + if !has_payload { + return Err(Error::invalid_input_source( + format!("Sparse structural {label} explicit encoding must require one payload") + .into(), + )); + } + Ok(Arc::from(decompressor)) } fn position_set_decoder( position_set: &pb21::SparsePositionSet, domain_len: u64, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result<(SparsePositionSetDecoder, u64)> { let cardinality = Self::position_cardinality(position_set, domain_len, label)?; let positions = position_set.positions.as_ref().ok_or_else(|| { @@ -2145,12 +2142,7 @@ impl SparseStructuralScheduler { } pb21::sparse_position_set::Positions::Explicit(compression) => { SparsePositionSetDecoder::Explicit { - decompressor: Self::create_position_decompressor( - compression, - label, - decompressors, - )?, - encoding: compression.clone(), + decompressor: Self::create_position_decompressor(compression, label)?, count: cardinality, domain_len, } @@ -2164,7 +2156,6 @@ impl SparseStructuralScheduler { validity_set: &pb21::SparseValiditySet, domain_len: u64, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result<(SparseValiditySetDecoder, u64)> { let meaning = Self::validity_meaning(validity_set, label)?; let position_set = validity_set.positions.as_ref().ok_or_else(|| { @@ -2172,8 +2163,7 @@ impl SparseStructuralScheduler { format!("Sparse structural {label} positions are required").into(), ) })?; - let (positions, cardinality) = - Self::position_set_decoder(position_set, domain_len, label, decompressors)?; + let (positions, cardinality) = Self::position_set_decoder(position_set, domain_len, label)?; Ok((SparseValiditySetDecoder { meaning, positions }, cardinality)) } @@ -2203,7 +2193,6 @@ impl SparseStructuralScheduler { count_set: &pb21::SparseCountSet, cardinality: u64, label: &str, - decompressors: &dyn DecompressionStrategy, ) -> Result { Self::count_buffer_count(count_set, cardinality, label)?; let counts = count_set.counts.as_ref().ok_or_else(|| { @@ -2219,12 +2208,7 @@ impl SparseStructuralScheduler { }, pb21::sparse_count_set::Counts::Explicit(compression) => { SparseCountSetDecoder::Explicit { - decompressor: Self::create_position_decompressor( - compression, - label, - decompressors, - )?, - encoding: compression.clone(), + decompressor: Self::create_position_decompressor(compression, label)?, count: cardinality, } } @@ -2335,7 +2319,6 @@ impl SparseStructuralScheduler { fn layer_decompressors( layer: &pb21::SparseStructuralLayer, - decompressors: &dyn DecompressionStrategy, ) -> Result { Ok(match Self::require_layer(layer)? { pb21::sparse_structural_layer::Layer::Validity(layer) => { @@ -2343,7 +2326,6 @@ impl SparseStructuralScheduler { Self::require_validity_set(&layer.validity, "validity")?, layer.num_slots, "validity positions", - decompressors, )?; SparseLayerDecompressors::Validity { num_slots: layer.num_slots, @@ -2357,19 +2339,16 @@ impl SparseStructuralScheduler { non_empty_positions, layer.num_slots, "list non-empty positions", - decompressors, )?; let counts = Self::count_set_decoder( Self::require_count_set(&layer.counts, "list counts")?, num_non_empty, "list counts", - decompressors, )?; let (validity, _) = Self::validity_set_decoder( Self::require_validity_set(&layer.validity, "list")?, layer.num_slots, "list validity positions", - decompressors, )?; SparseLayerDecompressors::List { num_slots: layer.num_slots, @@ -2384,7 +2363,6 @@ impl SparseStructuralScheduler { Self::require_validity_set(&layer.validity, "fixed-size-list")?, layer.num_slots, "fixed-size-list validity positions", - decompressors, )?; SparseLayerDecompressors::FixedSizeList { num_slots: layer.num_slots, @@ -2575,88 +2553,19 @@ impl SparseStructuralScheduler { Ok(()) } - fn validate_structural_buffer_headers( - encoding: &CompressiveEncoding, - data: &[u8], - label: &str, - ) -> Result<()> { - use pb21::compressive_encoding::Compression; - - match encoding.compression.as_ref() { - Some(Compression::General(general)) => { - Self::validate_general_buffer_header(general, data, label) - } - Some(Compression::Rle(rle)) => { - let values_size = u64::from_le_bytes( - data.get(..8) - .ok_or_else(|| { - Error::invalid_input_source( - format!( - "Sparse structural {label} RLE buffer is missing its header" - ) - .into(), - ) - })? - .try_into() - .map_err(|_| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE header is malformed").into(), - ) - })?, - ); - let values_size = usize_from_u64(values_size, "RLE values buffer size")?; - let values_end = 8_usize.checked_add(values_size).ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE values range overflows").into(), - ) - })?; - let values_data = data.get(8..values_end).ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE values buffer is truncated").into(), - ) - })?; - let lengths_data = data.get(values_end..).ok_or_else(|| { - Error::invalid_input_source( - format!("Sparse structural {label} RLE run-length buffer is missing") - .into(), - ) - })?; - Self::validate_general_child_buffer( - Self::require_encoding(&rle.values, "RLE values")?, - values_data, - "RLE values", - )?; - Self::validate_general_child_buffer( - Self::require_encoding(&rle.run_lengths, "RLE run lengths")?, - lengths_data, - "RLE run lengths", - ) - } - _ => Ok(()), - } - } - fn decode_u64_values( decompressor: &dyn BlockDecompressor, - encoding: &CompressiveEncoding, data: Bytes, num_values: u64, label: &str, ) -> Result> { - Self::validate_structural_buffer_headers(encoding, &data, label)?; - let decoded = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - decompressor.decompress(LanceBuffer::from_bytes(data, 1), num_values) - })) - .map_err(|_| { - Error::invalid_input_source( - format!("Sparse structural {label} decompression panicked").into(), - ) - })? - .map_err(|error| { - Error::invalid_input_source( - format!("Sparse structural {label} decompression failed: {error}").into(), - ) - })?; + let decoded = decompressor + .decompress(Some(LanceBuffer::from_bytes(data, 1)), num_values) + .map_err(|error| { + Error::invalid_input_source( + format!("Sparse structural {label} decompression failed: {error}").into(), + ) + })?; let fixed = decoded.as_fixed_width().ok_or_else(|| { Error::invalid_input_source( format!("Sparse structural {label} did not decode to fixed width data").into(), @@ -2718,14 +2627,12 @@ impl SparseStructuralScheduler { fn decode_explicit_positions( decompressor: &Arc, - encoding: &CompressiveEncoding, data: Bytes, num_positions: u64, num_slots: u64, label: &str, ) -> Result { - let deltas = - Self::decode_u64_values(decompressor.as_ref(), encoding, data, num_positions, label)?; + let deltas = Self::decode_u64_values(decompressor.as_ref(), data, num_positions, label)?; let mut positions = Vec::with_capacity(deltas.len()); let mut current = 0_u64; for (idx, delta) in deltas.into_iter().enumerate() { @@ -2771,12 +2678,10 @@ impl SparseStructuralScheduler { } SparsePositionSetDecoder::Explicit { decompressor, - encoding, count, domain_len, } => Self::decode_explicit_positions( decompressor, - encoding, Self::next_structural_buffer(buffers, label)?, *count, *domain_len, @@ -2808,12 +2713,10 @@ impl SparseStructuralScheduler { } SparseCountSetDecoder::Explicit { decompressor, - encoding, count, } => { let counts = Self::decode_u64_values( decompressor.as_ref(), - encoding, Self::next_structural_buffer(buffers, label)?, *count, label, @@ -4684,11 +4587,9 @@ fn slice_sparse_plan( mod tests { use std::sync::Mutex; - use crate::{ - compression::DefaultDecompressionStrategy, - encodings::physical::block::{CompressionConfig, CompressionScheme}, - testing::SimulatedScheduler, - }; + #[cfg(feature = "lz4")] + use crate::encodings::physical::block::{CompressionConfig, CompressionScheme}; + use crate::{compression::DefaultDecompressionStrategy, testing::SimulatedScheduler}; use super::*; @@ -4723,6 +4624,7 @@ mod tests { ) } + #[cfg(feature = "lz4")] fn general_lz4(values: CompressiveEncoding) -> CompressiveEncoding { ProtobufUtils21::wrapped(CompressionConfig::new(CompressionScheme::Lz4, None), values) .unwrap() diff --git a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs index 3d82b4a5ef6..fc3f3416cba 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/sparse/writer.rs @@ -10,7 +10,7 @@ use lance_core::{Error, Result, datatypes::Field, utils::bit::pad_bytes}; use crate::{ buffer::LanceBuffer, - compression::CompressionStrategy, + compression::{CompressionStrategy, compress_required_block}, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, decoder::PageEncoding, encoder::EncodedPage, @@ -24,7 +24,8 @@ use super::{ SparseValidityMeaning, SparseValiditySet, }; use crate::encodings::logical::primitive::{ - FILL_BYTE, MINIBLOCK_ALIGNMENT, miniblock::MiniBlockCompressed, + FILL_BYTE, MINIBLOCK_ALIGNMENT, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext}, }; #[derive(Clone, Copy, Default)] @@ -563,7 +564,8 @@ pub fn prepare_values( let num_values = data.num_values(); let compressor = compression_strategy.create_miniblock_compressor(field, &data)?; - let (compressed, value_compression) = compressor.compress(data)?; + let compression_context = MiniBlockCompressionContext::new(0, support_large_chunk, false); + let (compressed, value_compression) = compressor.compress(data, compression_context)?; let values = serialize_value_chunks(with_explicit_value_counts(compressed)?, support_large_chunk)?; Ok(PreparedSparseValues { @@ -586,8 +588,7 @@ fn encode_u64_values( }); block.compute_stat(); let field = Field::new_arrow("", arrow_schema::DataType::UInt64, false)?; - let (compressor, encoding) = compression_strategy.create_block_compressor(&field, &block)?; - Ok((compressor.compress(block)?, encoding)) + compress_required_block(compression_strategy, &field, block) } fn positions_to_deltas(positions: &[u64], label: &str) -> Result> { @@ -885,8 +886,9 @@ mod tests { use crate::{ constants::{ - PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, - STRUCTURAL_ENCODING_MINIBLOCK, STRUCTURAL_ENCODING_SPARSE, + COMPRESSION_META_KEY, PACKED_STRUCT_META_KEY, STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, + STRUCTURAL_ENCODING_SPARSE, }, data::FixedSizeListBlock, encoder::{ @@ -1543,6 +1545,62 @@ mod tests { check_round_trip_encoding_of_data(vec![array], &cases, sparse_metadata()).await; } + #[tokio::test] + async fn test_v2_3_sparse_variable_values_keep_legacy_flat_offsets() { + let num_values = 4_096; + let values = Arc::new(StringArray::from_iter_values((0..num_values).map( + |index| { + let length = [4_usize, 10, 5, 8][index % 4]; + format!("{index:04x}{}", "x".repeat(length - 4)) + }, + ))) as ArrayRef; + let array = + sparse_list_values( + num_values * 2, + 2, + values, + Arc::new(ArrowField::new("item", DataType::Utf8, true).with_metadata( + HashMap::from([(COMPRESSION_META_KEY.to_string(), "none".to_string())]), + )), + ); + let metadata = sparse_metadata(); + let pages = encode_pages(array.clone(), LanceFileVersion::V2_3, metadata.clone()) + .await + .unwrap(); + assert_eq!(pages.len(), 1); + let sparse = sparse_layout(&pages[0]); + assert_eq!(sparse.num_buffers, 1); + let Some(pb21::compressive_encoding::Compression::Variable(variable)) = sparse + .value_compression + .as_ref() + .and_then(|encoding| encoding.compression.as_ref()) + else { + panic!( + "expected sparse variable value compression, got {:?}", + sparse.value_compression + ); + }; + assert!(variable.values.is_none()); + assert!(matches!( + variable + .offsets + .as_deref() + .and_then(|offsets| offsets.compression.as_ref()), + Some(pb21::compressive_encoding::Compression::Flat(pb21::Flat { + bits_per_value: 32, + data: None, + })) + )); + + let cases = TestCases::default() + .with_min_file_version(LanceFileVersion::V2_3) + .with_max_file_version(LanceFileVersion::V2_3) + .with_page_sizes(vec![1]) + .with_range(1..17) + .with_indices(vec![0, 7, (num_values * 2 - 1) as u64]); + check_round_trip_encoding_of_data(vec![array], &cases, metadata).await; + } + #[tokio::test] async fn test_explicit_sparse_struct_with_constant_and_sparse_children() { let fields = Fields::from(vec![ diff --git a/rust/lance-encoding/src/encodings/physical.rs b/rust/lance-encoding/src/encodings/physical.rs index 0439c0216fb..46c77d271e2 100644 --- a/rust/lance-encoding/src/encodings/physical.rs +++ b/rust/lance-encoding/src/encodings/physical.rs @@ -1,14 +1,49 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors +use lance_core::{Error, Result}; + pub mod binary; #[cfg(feature = "bitpacking")] pub mod bitpacking; pub mod block; pub mod byte_stream_split; pub mod constant; +pub mod delta; +pub mod dictionary; pub mod fsst; pub mod general; pub mod packed; +pub mod range; pub mod rle; pub mod value; + +pub(crate) fn checked_vec_capacity( + num_values: u64, + bytes_per_value: usize, + label: &str, +) -> Result { + let capacity = usize::try_from(num_values) + .map_err(|_| Error::invalid_input(format!("{label} cardinality does not fit usize")))?; + let output_bytes = capacity + .checked_mul(bytes_per_value) + .ok_or_else(|| Error::invalid_input(format!("{label} byte length overflows usize")))?; + if output_bytes > isize::MAX as usize { + return Err(Error::invalid_input(format!( + "{label} byte length {output_bytes} exceeds isize::MAX" + ))); + } + Ok(capacity) +} + +pub(crate) fn try_vec_with_capacity(num_values: u64, label: &str) -> Result> { + let capacity = checked_vec_capacity(num_values, std::mem::size_of::(), label)?; + let output_bytes = capacity * std::mem::size_of::(); + let mut values = Vec::new(); + values.try_reserve_exact(capacity).map_err(|error| { + Error::invalid_input(format!( + "{label} could not reserve {capacity} values ({output_bytes} bytes): {error}" + )) + })?; + Ok(values) +} diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index d02cf2da693..34829710d3f 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -11,34 +11,45 @@ use arrow_array::OffsetSizeTrait; use byteorder::{ByteOrder, LittleEndian}; -use core::panic; +use prost::Message; use crate::compression::{ - BlockCompressor, BlockDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor, + BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, + VariablePerValueDecompressor, require_block_payload, }; use crate::buffer::LanceBuffer; -use crate::data::{BlockInfo, DataBlock, VariableWidthBlock}; +#[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::{ - MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, + MiniBlockCompressor, }; use crate::format::pb21::CompressiveEncoding; use crate::format::pb21::compressive_encoding::Compression; use crate::format::{ProtobufUtils21, pb21}; -use lance_core::utils::bit::pad_bytes_to; use lance_core::{Error, Result}; +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, } } } @@ -52,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] - .to_usize() - .unwrap(); - let end_in_orig = offsets[chunk.chunk_last_offset_in_orig_idx] + let start_in_orig = offsets[range.start_offset_index] .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. @@ -157,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; @@ -196,58 +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) + } + _ => 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); + } } - _ => panic!("Unsupported bits_per_offset={}", data.bits_per_offset), + 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) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { - DataBlock::VariableWidth(variable_width) => Ok(self.chunk_data(variable_width)), + 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", @@ -261,29 +836,82 @@ impl MiniBlockCompressor for BinaryMiniBlockEncoder { #[derive(Debug)] pub struct BinaryMiniBlockDecompressor { - bits_per_offset: u8, + layout: BinaryMiniBlockLayout, } -impl BinaryMiniBlockDecompressor { - pub fn new(bits_per_offset: u8) -> Self { - assert!(bits_per_offset == 32 || bits_per_offset == 64); - Self { bits_per_offset } +#[derive(Debug)] +enum BinaryMiniBlockLayout { + Legacy { + bits_per_offset: u8, + }, + Generic { + value_type: BlockValueType, + offsets: Box, + offsets_have_payload: bool, + validation: OffsetValidation, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OffsetValidation { + Endpoints, + Full, +} + +fn offset_validation(compression: &Compression) -> OffsetValidation { + match compression { + Compression::Constant(_) | Compression::Range(_) | Compression::Delta(_) => { + OffsetValidation::Endpoints + } + _ => OffsetValidation::Full, } +} - pub fn from_variable(variable: &pb21::Variable) -> Self { - if let Compression::Flat(flat) = variable +impl BinaryMiniBlockDecompressor { + pub fn from_variable(variable: &pb21::Variable) -> Result { + if variable.values.is_some() { + return Err(Error::invalid_input( + "Binary mini-block Variable values encoding must be absent", + )); + } + let offsets = variable .offsets .as_ref() - .unwrap() + .ok_or_else(|| Error::invalid_input("Variable encoding is missing offsets"))?; + let compression = offsets .compression .as_ref() - .unwrap() - { - Self { - bits_per_offset: flat.bits_per_value as u8, + .ok_or_else(|| Error::invalid_input("Variable offsets are missing compression"))?; + if let Compression::Flat(flat) = compression { + if flat.data.is_some() || !matches!(flat.bits_per_value, 32 | 64) { + return Err(Error::invalid_input(format!( + "Legacy Variable offsets require uncompressed 32 or 64-bit Flat encoding, got {} bits", + flat.bits_per_value + ))); } + Ok(Self { + layout: BinaryMiniBlockLayout::Legacy { + bits_per_offset: flat.bits_per_value as u8, + }, + }) } else { - panic!("Unsupported offsets compression: {:?}", variable.offsets); + let value_type = infer_block_value_type(offsets)?; + if !matches!(value_type, BlockValueType::UInt32 | BlockValueType::UInt64) { + return Err(Error::invalid_input(format!( + "Generic Variable offsets require u32 or u64 output, got {} bits", + value_type.bits_per_value() + ))); + } + let validation = offset_validation(compression); + let (offsets, offsets_have_payload) = create_block_decompressor(offsets, value_type)?; + Ok(Self { + layout: BinaryMiniBlockLayout::Generic { + value_type, + offsets, + offsets_have_payload, + validation, + }, + }) } } } @@ -294,53 +922,304 @@ impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { // it has so assertion can not be done here and the caller of `decompress` must ensure // `num_values` <= number of values in the chunk. fn decompress(&self, data: Vec, num_values: u64) -> Result { - assert_eq!(data.len(), 1); - let data = data.into_iter().next().unwrap(); - - if self.bits_per_offset == 64 { - // offset and at least one value - assert!(data.len() >= 16); - - let offsets_buffer = data.borrow_to_typed_slice::(); - let offsets = offsets_buffer.as_ref(); - - let result_offsets = offsets[0..(num_values + 1) as usize] - .iter() - .map(|offset| offset - offsets[0]) - .collect::>(); + match &self.layout { + BinaryMiniBlockLayout::Legacy { bits_per_offset } => { + decode_legacy_binary_miniblock(data, num_values, *bits_per_offset) + } + BinaryMiniBlockLayout::Generic { + value_type, + offsets, + offsets_have_payload, + validation, + } => decode_generic_binary_miniblock( + data, + num_values, + *value_type, + offsets.as_ref(), + *offsets_have_payload, + *validation, + ), + } + } +} +fn decode_legacy_binary_miniblock( + mut data: Vec, + num_values: u64, + bits_per_offset: u8, +) -> Result { + if data.len() != 1 { + return Err(Error::invalid_input(format!( + "Legacy Variable mini-block requires 1 buffer, got {}", + data.len() + ))); + } + let data = data.pop().expect("buffer count was checked"); + let num_offsets = num_values + .checked_add(1) + .ok_or_else(|| Error::invalid_input("Variable offset count overflows u64"))?; + let num_offsets = usize::try_from(num_offsets) + .map_err(|_| Error::invalid_input("Variable offset count does not fit usize"))?; + match bits_per_offset { + 32 => { + let offset_bytes = num_offsets.checked_mul(4).ok_or_else(|| { + Error::invalid_input("Variable offset byte length overflows usize") + })?; + if data.len() < offset_bytes { + return Err(Error::invalid_input(format!( + "Legacy Variable mini-block has {} bytes, shorter than {offset_bytes} offset bytes", + data.len() + ))); + } + let offsets_buffer = data.borrow_to_typed_slice::(); + let offsets = offsets_buffer + .get(..num_offsets) + .ok_or_else(|| Error::invalid_input("Legacy Variable offsets are truncated"))?; + let start = offsets[0]; + let mut result_offsets = Vec::with_capacity(num_offsets); + let mut previous = start; + for (index, offset) in offsets.iter().copied().enumerate() { + if offset < previous { + return Err(Error::invalid_input(format!( + "Legacy Variable offsets decrease at index {index}" + ))); + } + let relative = offset - start; + if relative > i32::MAX as u32 { + return Err(Error::invalid_input(format!( + "Legacy Variable relative offset {relative} exceeds i32::MAX" + ))); + } + result_offsets.push(relative); + previous = offset; + } + let start = usize::try_from(start) + .map_err(|_| Error::invalid_input("Variable data start does not fit usize"))?; + let end = usize::try_from(*offsets.last().expect("offsets are non-empty")) + .map_err(|_| Error::invalid_input("Variable data end does not fit usize"))?; + if start < offset_bytes || end < start || end > data.len() { + return Err(Error::invalid_input(format!( + "Legacy Variable data range {start}..{end} is invalid for {} bytes", + data.len() + ))); + } Ok(DataBlock::VariableWidth(VariableWidthBlock { - data: LanceBuffer::from( - data[offsets[0] as usize..offsets[num_values as usize] as usize].to_vec(), - ), + data: LanceBuffer::from(data[start..end].to_vec()), offsets: LanceBuffer::reinterpret_vec(result_offsets), - bits_per_offset: 64, + bits_per_offset, num_values, block_info: BlockInfo::new(), })) - } else { - // offset and at least one value - assert!(data.len() >= 8); - - let offsets_buffer = data.borrow_to_typed_slice::(); - let offsets = offsets_buffer.as_ref(); - - let result_offsets = offsets[0..(num_values + 1) as usize] - .iter() - .map(|offset| offset - offsets[0]) - .collect::>(); - + } + 64 => { + let offset_bytes = num_offsets.checked_mul(8).ok_or_else(|| { + Error::invalid_input("Variable offset byte length overflows usize") + })?; + if data.len() < offset_bytes { + return Err(Error::invalid_input(format!( + "Legacy Variable mini-block has {} bytes, shorter than {offset_bytes} offset bytes", + data.len() + ))); + } + let offsets_buffer = data.borrow_to_typed_slice::(); + let offsets = offsets_buffer + .get(..num_offsets) + .ok_or_else(|| Error::invalid_input("Legacy Variable offsets are truncated"))?; + let start = offsets[0]; + let mut result_offsets = Vec::with_capacity(num_offsets); + let mut previous = start; + for (index, offset) in offsets.iter().copied().enumerate() { + if offset < previous { + return Err(Error::invalid_input(format!( + "Legacy Variable offsets decrease at index {index}" + ))); + } + let relative = offset - start; + if relative > i64::MAX as u64 { + return Err(Error::invalid_input(format!( + "Legacy Variable relative offset {relative} exceeds i64::MAX" + ))); + } + result_offsets.push(relative); + previous = offset; + } + let start = usize::try_from(start) + .map_err(|_| Error::invalid_input("Variable data start does not fit usize"))?; + let end = usize::try_from(*offsets.last().expect("offsets are non-empty")) + .map_err(|_| Error::invalid_input("Variable data end does not fit usize"))?; + if start < offset_bytes || end < start || end > data.len() { + return Err(Error::invalid_input(format!( + "Legacy Variable data range {start}..{end} is invalid for {} bytes", + data.len() + ))); + } Ok(DataBlock::VariableWidth(VariableWidthBlock { - data: LanceBuffer::from( - data[offsets[0] as usize..offsets[num_values as usize] as usize].to_vec(), - ), + data: LanceBuffer::from(data[start..end].to_vec()), offsets: LanceBuffer::reinterpret_vec(result_offsets), - bits_per_offset: 32, + bits_per_offset, num_values, block_info: BlockInfo::new(), })) } + _ => Err(Error::invalid_input(format!( + "Legacy Variable offsets require 32 or 64 bits, got {bits_per_offset}" + ))), + } +} + +fn decode_generic_binary_miniblock( + mut data: Vec, + num_values: u64, + value_type: BlockValueType, + offsets_decoder: &dyn BlockDecompressor, + offsets_have_payload: bool, + validation: OffsetValidation, +) -> Result { + let num_offsets = num_values + .checked_add(1) + .ok_or_else(|| Error::invalid_input("Variable offset count overflows u64"))?; + let expected_buffers = 1 + usize::from(offsets_have_payload); + if data.len() != expected_buffers { + return Err(Error::invalid_input(format!( + "Generic Variable mini-block requires {expected_buffers} buffers, got {}", + data.len() + ))); } + let values = data.pop().expect("buffer count was checked"); + let offset_payload = if offsets_have_payload { + Some(data.pop().expect("buffer count was checked")) + } else { + None + }; + let offsets = offsets_decoder.decompress(offset_payload, num_offsets)?; + let DataBlock::FixedWidth(offsets) = offsets else { + return Err(Error::invalid_input( + "Generic Variable offsets decoded to a non fixed-width block", + )); + }; + if offsets.num_values != num_offsets || offsets.bits_per_value != value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Generic Variable offsets decoded {} {}-bit values, expected {num_offsets} {}-bit values", + offsets.num_values, + offsets.bits_per_value, + value_type.bits_per_value() + ))); + } + validate_decoded_offsets(&offsets, value_type, values.len(), validation)?; + Ok(DataBlock::VariableWidth(VariableWidthBlock { + data: values, + offsets: offsets.data, + bits_per_offset: value_type.bits_per_value() as u8, + num_values, + block_info: BlockInfo::new(), + })) +} + +fn validate_decoded_offsets( + offsets: &FixedWidthDataBlock, + value_type: BlockValueType, + data_len: usize, + validation: OffsetValidation, +) -> Result<()> { + let signed_max = match value_type { + BlockValueType::UInt32 => i32::MAX as u64, + BlockValueType::UInt64 => i64::MAX as u64, + _ => { + return Err(Error::invalid_input( + "Generic Variable offsets require u32 or u64 values", + )); + } + }; + if validation == OffsetValidation::Endpoints { + let (first, last) = match value_type { + BlockValueType::UInt32 => { + let offsets = offsets.data.borrow_to_typed_slice::(); + let first = offsets.first().copied().ok_or_else(|| { + Error::invalid_input("Generic Variable offsets cannot be empty") + })?; + let last = offsets.last().copied().expect("first offset was present"); + (u64::from(first), u64::from(last)) + } + BlockValueType::UInt64 => { + let offsets = offsets.data.borrow_to_typed_slice::(); + let first = offsets.first().copied().ok_or_else(|| { + Error::invalid_input("Generic Variable offsets cannot be empty") + })?; + let last = offsets.last().copied().expect("first offset was present"); + (first, last) + } + _ => unreachable!("generic offset type was validated"), + }; + if first != 0 { + return Err(Error::invalid_input(format!( + "Generic Variable offsets must start at zero, got {first}" + ))); + } + if last > signed_max { + return Err(Error::invalid_input(format!( + "Generic Variable offset {last} exceeds the Arrow signed offset range" + ))); + } + if last != data_len as u64 { + return Err(Error::invalid_input(format!( + "Generic Variable final offset {last} does not equal {data_len} value bytes" + ))); + } + return Ok(()); + } + let mut previous = None; + let mut observe = |index: usize, offset: u64| -> Result<()> { + if index == 0 && offset != 0 { + return Err(Error::invalid_input(format!( + "Generic Variable offsets must start at zero, got {offset}" + ))); + } + if previous.is_some_and(|previous| offset < previous) { + return Err(Error::invalid_input(format!( + "Generic Variable offsets decrease at index {index}" + ))); + } + if offset > signed_max { + return Err(Error::invalid_input(format!( + "Generic Variable offset {offset} exceeds the Arrow signed offset range" + ))); + } + previous = Some(offset); + Ok(()) + }; + match value_type { + BlockValueType::UInt32 => { + for (index, offset) in offsets + .data + .borrow_to_typed_slice::() + .iter() + .copied() + .enumerate() + { + observe(index, u64::from(offset))?; + } + } + BlockValueType::UInt64 => { + for (index, offset) in offsets + .data + .borrow_to_typed_slice::() + .iter() + .copied() + .enumerate() + { + observe(index, offset)?; + } + } + _ => unreachable!("generic offset type was validated"), + } + let final_offset = + previous.ok_or_else(|| Error::invalid_input("Generic Variable offsets cannot be empty"))?; + if final_offset != data_len as u64 { + return Err(Error::invalid_input(format!( + "Generic Variable final offset {final_offset} does not equal {data_len} value bytes" + ))); + } + Ok(()) } /// Most basic encoding for variable-width data which does no compression at all @@ -356,7 +1235,7 @@ impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { pub struct VariableEncoder {} impl BlockCompressor for VariableEncoder { - fn compress(&self, mut data: DataBlock) -> Result { + fn compress(&self, mut data: DataBlock) -> Result> { match data { DataBlock::VariableWidth(ref mut variable_width_data) => { match variable_width_data.bits_per_offset { @@ -408,18 +1287,17 @@ impl BlockCompressor for VariableEncoder { output.extend_from_slice(&variable_width_data.data); Ok(LanceBuffer::from(output)) } - _ => { - panic!( - "BinaryBlockEncoder does not work with {} bits per offset VariableWidth DataBlock.", - variable_width_data.bits_per_offset - ); - } + _ => Err(Error::invalid_input(format!( + "BinaryBlockEncoder does not support {}-bit offsets", + variable_width_data.bits_per_offset + ))), } } - _ => { - panic!("BinaryBlockEncoder can only work with Variable Width DataBlock."); - } + _ => Err(Error::invalid_input( + "BinaryBlockEncoder requires a variable-width block", + )), } + .map(Some) } } @@ -450,7 +1328,8 @@ impl VariablePerValueDecompressor for VariableDecoder { pub struct BinaryBlockDecompressor {} impl BlockDecompressor for BinaryBlockDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Binary block")?; // In older (not quite stable) versions we stored the bits per offset as a single byte and then the num_values // as four bytes. However, this led to alignment problems and was wasteful since we already store the num_values // in higher layers. @@ -526,6 +1405,7 @@ impl BlockDecompressor for BinaryBlockDecompressor { #[cfg(test)] mod tests { + use super::*; use arrow_array::{ ArrayRef, StringArray, builder::{LargeStringBuilder, StringBuilder}, @@ -534,8 +1414,8 @@ mod tests { use crate::{ constants::{ - COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, - STRUCTURAL_ENCODING_MINIBLOCK, + COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY, STRUCTURAL_ENCODING_FULLZIP, + STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, }, testing::check_specific_random, }; @@ -543,6 +1423,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, @@ -550,6 +1431,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; + } + #[test_log::test(tokio::test)] async fn test_utf8_binary() { let field = Field::new("", DataType::Utf8, false); @@ -863,7 +2146,9 @@ mod tests { // Test case 1: u32 offsets { let decompressor = BinaryMiniBlockDecompressor { - bits_per_offset: 32, + layout: super::BinaryMiniBlockLayout::Legacy { + bits_per_offset: 32, + }, }; // Create test data with u32 offsets @@ -919,7 +2204,9 @@ mod tests { // Test case 2: u64 offsets { let decompressor = BinaryMiniBlockDecompressor { - bits_per_offset: 64, + layout: super::BinaryMiniBlockLayout::Legacy { + bits_per_offset: 64, + }, }; // Create test data with u64 offsets diff --git a/rust/lance-encoding/src/encodings/physical/binary/offsets.rs b/rust/lance-encoding/src/encodings/physical/binary/offsets.rs new file mode 100644 index 00000000000..5212d7f0377 --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/binary/offsets.rs @@ -0,0 +1,1449 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Adopter-private selection for chunk-local variable-width offsets. + +use std::sync::Arc; + +use prost::Message; + +use crate::{ + buffer::LanceBuffer, + compression::block::{ + MAX_DICTIONARY_ITEMS, fixed_from_u64_values, validate_fixed_payload_len, + visit_unsigned_values, + }, + compression::{BlockCompressor, BlockValueType}, + compression_config::CompressionFieldParams, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + encodings::physical::{ + delta::encode_deltas, + dictionary::{self, DictionaryBlockCompressor}, + rle::{self, BlockRleCompressor}, + }, + format::{ProtobufUtils21, pb21::CompressiveEncoding}, +}; +use lance_core::{Error, Result}; + +mod selector; +mod statistics; + +use selector::{ + Candidate, constant_selection, direct_candidates, estimate_payload, flat_selection, + range_selection, +}; +use statistics::{BoundedDistinctStatsBuilder, SequenceStats, SequenceStatsBuilder}; + +const GENERAL_SAMPLE_BYTES: usize = 64 * 1024; + +#[derive(Debug, Clone, Copy, Default)] +pub(super) struct BlockCost { + buffer_overhead_bytes: u64, + alignment: u64, +} + +impl BlockCost { + pub(super) fn new(buffer_overhead_bytes: u64, alignment: u64) -> Self { + Self { + buffer_overhead_bytes, + alignment: alignment.max(1), + } + } + + fn payload_wire_bytes(self, has_payload: bool, payload_bytes: u64) -> u64 { + align_wire_bytes(payload_bytes, self.alignment.max(1)).saturating_add(if has_payload { + self.buffer_overhead_bytes + } else { + 0 + }) + } +} + +fn align_wire_bytes(bytes: u64, alignment: u64) -> u64 { + bytes + .checked_add(alignment - 1) + .map(|padded| padded / alignment * alignment) + .unwrap_or(u64::MAX) +} + +#[derive(Debug)] +struct SelectedBlockCompressor { + compressor: Box, + encoding: CompressiveEncoding, + has_payload: bool, +} + +impl SelectedBlockCompressor { + fn new( + compressor: Box, + encoding: CompressiveEncoding, + has_payload: bool, + ) -> Self { + Self { + compressor, + encoding, + has_payload, + } + } +} + +#[derive(Debug, Clone, Copy)] +enum OffsetMetadataPattern { + Constant(u64), + Range { start: u64, step: u64 }, +} + +fn probe_offset_metadata( + data: &FixedWidthDataBlock, + ranges: &[std::ops::Range], +) -> Result> { + let value_type = BlockValueType::from_bits(data.bits_per_value)?; + validate_fixed_payload_len( + &data.data, + value_type, + data.num_values, + "Block family input", + )?; + match value_type { + BlockValueType::UInt32 => { + let values = data.data.borrow_to_typed_view::(); + probe_typed_offset_metadata(values.as_ref(), ranges, u64::from) + } + BlockValueType::UInt64 => { + let values = data.data.borrow_to_typed_view::(); + probe_typed_offset_metadata(values.as_ref(), ranges, |value| value) + } + _ => Err(Error::invalid_input(format!( + "Generic block family selection only supports u32 or u64, got {} bits", + value_type.bits_per_value() + ))), + } +} + +fn probe_typed_offset_metadata( + values: &[T], + ranges: &[std::ops::Range], + to_u64: impl Fn(T) -> u64 + Copy, +) -> Result> { + let mut common_step = None; + let mut all_constant = true; + let mut range_possible = true; + + for (index, range) in ranges.iter().enumerate() { + let member = values.get(range.clone()).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} range {}..{} exceeds {} values", + range.start, + range.end, + values.len() + )) + })?; + let base = + member.first().copied().map(to_u64).ok_or_else(|| { + Error::invalid_input("Variable-width offset chunks cannot be empty") + })?; + let mut previous = base; + for value in member.iter().copied().skip(1).map(to_u64) { + let delta = value.checked_sub(previous).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} must be non-decreasing" + )) + })?; + all_constant &= delta == 0; + match common_step { + Some(step) => range_possible &= delta == step, + None => common_step = Some(delta), + } + previous = value; + if !all_constant && !range_possible { + return Ok(None); + } + } + range_possible &= member.len() >= 2; + } + + if all_constant { + return Ok(Some(OffsetMetadataPattern::Constant(0))); + } + Ok(match (range_possible, common_step) { + (true, Some(step)) if step > 0 => Some(OffsetMetadataPattern::Range { start: 0, step }), + _ => None, + }) +} + +#[derive(Debug, Clone)] +pub(super) struct OffsetSequenceAnalysis { + pub(super) values: SequenceStats, + pub(super) deltas: Option, +} + +#[derive(Debug)] +pub(super) struct OffsetRunAnalysis { + pub(super) values: SequenceStats, + pub(super) lengths: SequenceStats, +} + +#[derive(Debug)] +pub(super) struct OffsetFamilyAnalysis { + pub(super) members: Vec, + pub(super) runs: Option>, + pub(super) dictionary_items: Option>, +} + +pub(super) fn analyze_offset_members( + data: &FixedWidthDataBlock, + ranges: &[std::ops::Range], + collect_sample: bool, +) -> Result { + let value_type = BlockValueType::from_bits(data.bits_per_value)?; + match value_type { + BlockValueType::UInt32 => { + let values = data.data.borrow_to_typed_view::(); + analyze_typed_offsets::<_, false>( + values.as_ref(), + ranges, + value_type, + collect_sample, + u64::from, + ) + } + BlockValueType::UInt64 => { + let values = data.data.borrow_to_typed_view::(); + analyze_typed_offsets::<_, true>( + values.as_ref(), + ranges, + value_type, + collect_sample, + |value| value, + ) + } + _ => Err(Error::invalid_input(format!( + "Generic block family selection only supports u32 or u64, got {} bits", + value_type.bits_per_value() + ))), + } +} + +fn analyze_typed_offsets( + values: &[T], + ranges: &[std::ops::Range], + value_type: BlockValueType, + collect_sample: bool, + to_u64: impl Fn(T) -> u64 + Copy, +) -> Result { + if !collect_sample { + return analyze_typed_offsets_without_samples::<_, COLLECT_DICTIONARY>( + values, ranges, value_type, to_u64, + ); + } + if ranges.is_empty() { + return Err(Error::invalid_input( + "Variable-width offsets require at least one chunk", + )); + } + let sample_limit = if collect_sample { + GENERAL_SAMPLE_BYTES / ranges.len() + } else { + 0 + }; + let mut dictionary_items = COLLECT_DICTIONARY.then(BoundedDistinctStatsBuilder::new); + let members = ranges + .iter() + .enumerate() + .map(|(index, range)| { + let member = values.get(range.clone()).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} range {}..{} exceeds {} values", + range.start, + range.end, + values.len() + )) + })?; + let base = member.first().copied().map(to_u64).ok_or_else(|| { + Error::invalid_input("Variable-width offset chunks cannot be empty") + })?; + let mut value_stats = SequenceStatsBuilder::with_sample_limit(value_type, sample_limit); + let mut delta_stats = SequenceStatsBuilder::with_sample_limit(value_type, sample_limit); + let mut previous = None; + for value in member.iter().copied().map(to_u64) { + let value = value.checked_sub(base).ok_or_else(|| { + Error::invalid_input( + "Block family member is smaller than its normalization base", + ) + })?; + value_stats.push(value); + if COLLECT_DICTIONARY { + dictionary_items + .as_mut() + .expect("dictionary collection is enabled") + .push(value); + } + if let Some(previous) = previous { + delta_stats.push(value.checked_sub(previous).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} must be non-decreasing" + )) + })?); + } + previous = Some(value); + } + let value_stats = value_stats.finish(); + if !value_stats.is_non_decreasing { + return Err(Error::invalid_input(format!( + "Block family member {index} must be non-decreasing" + ))); + } + let deltas = (value_stats.len >= 2).then(|| delta_stats.finish()); + Ok(OffsetSequenceAnalysis { + values: value_stats, + deltas, + }) + }) + .collect::>>()?; + let runs = analyze_offset_runs(values, ranges, value_type, collect_sample, to_u64, &members)?; + Ok(OffsetFamilyAnalysis { + members, + runs, + dictionary_items: dictionary_items.and_then(|items| items.finish()), + }) +} + +fn analyze_typed_offsets_without_samples( + values: &[T], + ranges: &[std::ops::Range], + value_type: BlockValueType, + to_u64: impl Fn(T) -> u64 + Copy, +) -> Result { + if ranges.is_empty() { + return Err(Error::invalid_input( + "Variable-width offsets require at least one chunk", + )); + } + let mut dictionary_items = COLLECT_DICTIONARY.then(BoundedDistinctStatsBuilder::new); + let members = ranges + .iter() + .enumerate() + .map(|(index, range)| { + let member = values.get(range.clone()).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} range {}..{} exceeds {} values", + range.start, + range.end, + values.len() + )) + })?; + let first_raw = member.first().copied().map(to_u64).ok_or_else(|| { + Error::invalid_input("Variable-width offset chunks cannot be empty") + })?; + let base = first_raw; + let first = 0_u64; + let mut previous = first; + let mut max = first; + let mut run_count = 1_u64; + + let mut delta_len = 0_u64; + let mut delta_first = None; + let mut delta_previous = None; + let mut delta_min = u64::MAX; + let mut delta_max = 0_u64; + let mut deltas_non_decreasing = true; + let mut delta_step = None; + let mut deltas_are_arithmetic = true; + let mut delta_run_count = 0_u64; + + if COLLECT_DICTIONARY { + dictionary_items + .as_mut() + .expect("dictionary collection is enabled") + .push(first); + } + for value in member.iter().copied().skip(1).map(to_u64) { + let value = value.checked_sub(base).ok_or_else(|| { + Error::invalid_input( + "Block family member is smaller than its normalization base", + ) + })?; + let delta = value.checked_sub(previous).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} must be non-decreasing" + )) + })?; + if value != previous { + run_count += 1; + } + previous = value; + max = value; + if COLLECT_DICTIONARY { + dictionary_items + .as_mut() + .expect("dictionary collection is enabled") + .push(value); + } + + if delta_first.is_none() { + delta_first = Some(delta); + delta_run_count = 1; + } + if let Some(previous_delta) = delta_previous { + deltas_non_decreasing &= delta >= previous_delta; + if delta != previous_delta { + delta_run_count += 1; + } + let difference = delta.checked_sub(previous_delta); + match (delta_step, difference) { + (None, Some(difference)) if delta_len == 1 => { + delta_step = Some(difference); + } + (Some(step), Some(difference)) if step == difference => {} + _ => deltas_are_arithmetic = false, + } + } + delta_previous = Some(delta); + delta_min = delta_min.min(delta); + delta_max = delta_max.max(delta); + delta_len += 1; + } + + let deltas = (delta_len > 0).then(|| SequenceStats { + value_type, + len: delta_len, + first: delta_first, + min: delta_min, + max: delta_max, + is_non_decreasing: deltas_non_decreasing, + arithmetic_step: (delta_len >= 2 && deltas_are_arithmetic) + .then_some(delta_step) + .flatten(), + run_count: delta_run_count, + sample: None, + }); + Ok(OffsetSequenceAnalysis { + values: SequenceStats { + value_type, + len: member.len() as u64, + first: Some(first), + min: first, + max, + is_non_decreasing: true, + arithmetic_step: None, + run_count, + sample: None, + }, + deltas, + }) + }) + .collect::>>()?; + let runs = analyze_offset_runs(values, ranges, value_type, false, to_u64, &members)?; + Ok(OffsetFamilyAnalysis { + members, + runs, + dictionary_items: dictionary_items.and_then(BoundedDistinctStatsBuilder::finish), + }) +} + +fn analyze_offset_runs( + values: &[T], + ranges: &[std::ops::Range], + value_type: BlockValueType, + collect_sample: bool, + to_u64: impl Fn(T) -> u64 + Copy, + members: &[OffsetSequenceAnalysis], +) -> Result>> { + let (total_values, total_runs) = + members + .iter() + .try_fold((0_u64, 0_u64), |(values, runs), member| { + Ok::<_, Error>(( + values.checked_add(member.values.len).ok_or_else(|| { + Error::invalid_input("Offset family cardinality overflows u64") + })?, + runs.checked_add(member.values.run_count).ok_or_else(|| { + Error::invalid_input("Offset family run count overflows u64") + })?, + )) + })?; + if total_runs == 0 || total_runs >= total_values { + return Ok(None); + } + + let sample_limit = if collect_sample { + GENERAL_SAMPLE_BYTES / ranges.len() + } else { + 0 + }; + ranges + .iter() + .enumerate() + .map(|(index, range)| { + let member = values.get(range.clone()).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} range {}..{} exceeds {} values", + range.start, + range.end, + values.len() + )) + })?; + let base = member.first().copied().map(to_u64).ok_or_else(|| { + Error::invalid_input("Variable-width offset chunks cannot be empty") + })?; + let mut run_values = SequenceStatsBuilder::with_sample_limit(value_type, sample_limit); + let mut run_lengths = SequenceStatsBuilder::without_sample(BlockValueType::UInt64); + let mut current = None; + let mut length = 0_u64; + for value in member.iter().copied().map(to_u64) { + let value = value.checked_sub(base).ok_or_else(|| { + Error::invalid_input( + "Block family member is smaller than its normalization base", + ) + })?; + match current { + Some(run_value) if run_value == value => { + length = length.checked_add(1).ok_or_else(|| { + Error::invalid_input("Offset run length overflows u64") + })?; + } + Some(run_value) => { + run_values.push(run_value); + run_lengths.push(length); + current = Some(value); + length = 1; + } + None => { + current = Some(value); + length = 1; + } + } + } + if let Some(run_value) = current { + run_values.push(run_value); + run_lengths.push(length); + } + Ok(OffsetRunAnalysis { + values: run_values.finish(), + lengths: run_lengths.finish(), + }) + }) + .collect::>>() + .map(Some) +} + +pub(super) fn offset_member_base( + block: &FixedWidthDataBlock, + value_type: BlockValueType, +) -> Result { + if block.num_values == 0 { + return Err(Error::invalid_input( + "Variable-width offset chunks cannot be empty", + )); + } + Ok(match value_type { + BlockValueType::UInt8 => u64::from(block.data[0]), + BlockValueType::UInt16 => u64::from(block.data.borrow_to_typed_view::().as_ref()[0]), + BlockValueType::UInt32 => u64::from(block.data.borrow_to_typed_view::().as_ref()[0]), + BlockValueType::UInt64 => block.data.borrow_to_typed_view::().as_ref()[0], + }) +} + +pub(super) fn normalize_offset_member(block: FixedWidthDataBlock) -> Result { + let value_type = BlockValueType::from_bits(block.bits_per_value)?; + validate_fixed_payload_len( + &block.data, + value_type, + block.num_values, + "Block family member", + )?; + let base = offset_member_base(&block, value_type)?; + if base == 0 { + return Ok(block); + } + let capacity = usize::try_from(block.num_values) + .map_err(|_| Error::invalid_input("Block family cardinality does not fit usize"))?; + let mut normalized = Vec::with_capacity(capacity); + visit_unsigned_values(&block, value_type, |value| { + normalized.push(value.checked_sub(base).ok_or_else(|| { + Error::invalid_input("Block family member is smaller than its normalization base") + })?); + Ok(()) + })?; + fixed_from_u64_values(&normalized, value_type, "Block family member") +} + +#[derive(Debug)] +struct NormalizedOffsetCompressor { + child: Box, +} + +impl NormalizedOffsetCompressor { + fn new(child: Box) -> Self { + Self { child } + } +} + +impl BlockCompressor for NormalizedOffsetCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Offset compression requires fixed-width data", + )); + }; + self.child + .compress(DataBlock::FixedWidth(normalize_offset_member(data)?)) + } +} + +#[derive(Debug)] +struct DeltaOffsetCompressor { + child: Box, +} + +impl DeltaOffsetCompressor { + fn new(child: Box) -> Self { + Self { child } + } +} + +impl BlockCompressor for DeltaOffsetCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Delta offset compression requires fixed-width data", + )); + }; + let value_type = BlockValueType::from_bits(data.bits_per_value)?; + let base = offset_member_base(&data, value_type)?; + let deltas = encode_deltas(data, base)?; + self.child.compress(DataBlock::FixedWidth(deltas)) + } +} + +#[derive(Debug, Clone)] +pub(super) struct OffsetFamilyMembers { + data: FixedWidthDataBlock, + ranges: Arc<[std::ops::Range]>, +} + +impl OffsetFamilyMembers { + fn into_blocks(self) -> Result> { + let bytes_per_value = usize::try_from(self.data.bits_per_value / 8) + .map_err(|_| Error::invalid_input("Block family value width does not fit usize"))?; + self.ranges + .iter() + .map(|range| { + let byte_start = range + .start + .checked_mul(bytes_per_value) + .ok_or_else(|| Error::invalid_input("Block family start overflows usize"))?; + let byte_len = range + .len() + .checked_mul(bytes_per_value) + .ok_or_else(|| Error::invalid_input("Block family length overflows usize"))?; + Ok(FixedWidthDataBlock { + data: self.data.data.slice_with_length(byte_start, byte_len), + bits_per_value: self.data.bits_per_value, + num_values: range.len() as u64, + block_info: BlockInfo::default(), + }) + }) + .collect() + } +} + +/// One reusable concrete compressor for all independently framed chunks. +#[derive(Debug)] +pub(super) struct OffsetFamilyCompressor { + compressor: Box, + encoding: CompressiveEncoding, + has_payload: bool, + estimated_payload_bytes: Arc<[u64]>, + members: OffsetFamilyMembers, +} + +impl OffsetFamilyCompressor { + pub(super) fn encoding(&self) -> &CompressiveEncoding { + &self.encoding + } + + pub(super) fn has_payload(&self) -> bool { + self.has_payload + } + + pub(super) fn estimated_payload_bytes(&self) -> &[u64] { + &self.estimated_payload_bytes + } + + pub(super) fn compress_members(self) -> Result>> { + let has_payload = self.has_payload; + if !has_payload { + return Ok(vec![None; self.members.ranges.len()]); + } + self.members + .into_blocks()? + .into_iter() + .zip(self.estimated_payload_bytes.iter().copied()) + .enumerate() + .map(|(index, (data, estimated_payload_bytes))| { + let payload = self.compressor.compress(DataBlock::FixedWidth(data))?; + if payload.is_some() != has_payload { + return Err(Error::internal(format!( + "Block family compressor payload presence changed for member {index}" + ))); + } + let payload_bytes = payload.as_ref().map_or(0, LanceBuffer::len) as u64; + if payload_bytes != estimated_payload_bytes { + return Err(Error::internal(format!( + "Block family member {index} estimated {estimated_payload_bytes} bytes but built {payload_bytes} bytes", + ))); + } + Ok(payload) + }) + .collect() + } +} + +#[derive(Debug)] +pub(super) struct OffsetFamilyCandidate { + selection: SelectedBlockCompressor, + estimated_payload_bytes: Vec, + estimated_wire_bytes: u64, + transform_depth: u8, + decode_cpu_rank: u8, + stable_rank: u8, +} + +impl OffsetFamilyCandidate { + fn is_better_than(&self, other: &Self) -> bool { + ( + self.estimated_wire_bytes, + self.transform_depth, + self.decode_cpu_rank, + self.stable_rank, + ) < ( + other.estimated_wire_bytes, + other.transform_depth, + other.decode_cpu_rank, + other.stable_rank, + ) + } + + fn finish(self, members: OffsetFamilyMembers) -> OffsetFamilyCompressor { + OffsetFamilyCompressor { + compressor: self.selection.compressor, + encoding: self.selection.encoding, + has_payload: self.selection.has_payload, + estimated_payload_bytes: Arc::from(self.estimated_payload_bytes), + members, + } + } + + fn normalize_members(mut self) -> Self { + self.selection.compressor = + Box::new(NormalizedOffsetCompressor::new(self.selection.compressor)); + self + } +} + +/// Builds a Delta(Flat) family after a bounded preflight has proven that a +/// complete selector pass cannot change the winner. +pub(super) fn select_delta_flat_offsets( + data: FixedWidthDataBlock, + member_ranges: Vec>, +) -> Result { + if member_ranges.is_empty() { + return Err(Error::invalid_input( + "Block sequence family requires at least one block", + )); + } + let value_type = BlockValueType::from_bits(data.bits_per_value)?; + validate_fixed_payload_len( + &data.data, + value_type, + data.num_values, + "Block family input", + )?; + let values_len = usize::try_from(data.num_values) + .map_err(|_| Error::invalid_input("Block family cardinality does not fit usize"))?; + let estimated_payload_bytes = member_ranges + .iter() + .enumerate() + .map(|(index, range)| { + if range.end > values_len || range.len() < 2 { + return Err(Error::invalid_input(format!( + "Delta block family member {index} range {}..{} is invalid for {values_len} values", + range.start, range.end + ))); + } + (range.len() as u64 - 1) + .checked_mul(value_type.bytes_per_value() as u64) + .ok_or_else(|| Error::invalid_input("Delta family payload size overflows u64")) + }) + .collect::>>()?; + validate_monotonic_offset_members(&data, &member_ranges, value_type)?; + + let child = flat_selection(value_type); + let child_has_payload = child.has_payload; + let encoding = ProtobufUtils21::delta(value_type.bits_per_value(), 0, child.encoding); + let selection = SelectedBlockCompressor::new( + Box::new(DeltaOffsetCompressor::new(child.compressor)), + encoding, + child_has_payload, + ); + Ok(OffsetFamilyCompressor { + compressor: selection.compressor, + encoding: selection.encoding, + has_payload: selection.has_payload, + estimated_payload_bytes: Arc::from(estimated_payload_bytes), + members: OffsetFamilyMembers { + data, + ranges: Arc::from(member_ranges), + }, + }) +} + +pub(super) fn validate_monotonic_offset_members( + data: &FixedWidthDataBlock, + ranges: &[std::ops::Range], + value_type: BlockValueType, +) -> Result<()> { + macro_rules! validate { + ($ty:ty) => {{ + let values = data.data.borrow_to_typed_view::<$ty>(); + for (index, range) in ranges.iter().enumerate() { + let member = values.get(range.clone()).ok_or_else(|| { + Error::invalid_input(format!( + "Block family member {index} range {}..{} exceeds {} values", + range.start, + range.end, + values.len() + )) + })?; + if member.windows(2).any(|pair| pair[1] < pair[0]) { + return Err(Error::invalid_input(format!( + "Block family member {index} must be non-decreasing" + ))); + } + } + }}; + } + match value_type { + BlockValueType::UInt32 => validate!(u32), + BlockValueType::UInt64 => validate!(u64), + _ => { + return Err(Error::invalid_input(format!( + "Generic block family selection only supports u32 or u64, got {} bits", + value_type.bits_per_value() + ))); + } + } + Ok(()) +} + +/// Selects one reusable concrete compressor for independent offset chunks. +pub(super) fn select_offset_family( + data: FixedWidthDataBlock, + member_ranges: Vec>, + field_params: &CompressionFieldParams, + cost: BlockCost, +) -> Result { + if member_ranges.is_empty() { + return Err(Error::invalid_input( + "Block sequence family requires at least one block", + )); + } + let value_type = BlockValueType::from_bits(data.bits_per_value)?; + if let Some(pattern) = probe_offset_metadata(&data, &member_ranges)? { + let candidate = match pattern { + OffsetMetadataPattern::Constant(value) => offset_metadata_candidate( + constant_selection(value_type, Some(value)), + member_ranges.len(), + cost, + 0, + ), + OffsetMetadataPattern::Range { start, step } => offset_metadata_candidate( + range_selection(value_type, start, step), + member_ranges.len(), + cost, + 1, + ), + } + .normalize_members(); + return Ok(candidate.finish(OffsetFamilyMembers { + data, + ranges: Arc::from(member_ranges), + })); + } + let collect_sample = !matches!(field_params.compression.as_deref(), Some("none" | "fsst")); + let analysis = analyze_offset_members(&data, &member_ranges, collect_sample)?; + let members = OffsetFamilyMembers { + data, + ranges: Arc::from(member_ranges), + }; + let value_stats = analysis + .members + .iter() + .map(|analysis| &analysis.values) + .collect::>(); + + let mut candidates = offset_direct_candidates(&value_stats, field_params, cost, 0, true)? + .into_iter() + .map(OffsetFamilyCandidate::normalize_members) + .collect::>(); + if let Some(candidate) = offset_rle_candidate(&analysis, field_params, cost)? { + candidates.push(candidate.normalize_members()); + } + if let Some(candidate) = offset_dictionary_candidate(&analysis, field_params, cost)? { + candidates.push(candidate.normalize_members()); + } + if value_stats.iter().all(|stats| stats.len >= 2) + && let Some(base) = common_first(&value_stats) + && analysis + .members + .iter() + .all(|analysis| analysis.deltas.is_some()) + { + let delta_stats = analysis + .members + .iter() + .map(|analysis| { + analysis + .deltas + .as_ref() + .expect("delta presence was checked") + }) + .collect::>(); + let child = offset_leaf_candidate(&delta_stats, field_params, cost)?; + let estimated_payload_bytes = child.estimated_payload_bytes.clone(); + let child_has_payload = child.selection.has_payload; + let encoding = + ProtobufUtils21::delta(value_type.bits_per_value(), base, child.selection.encoding); + let selection = SelectedBlockCompressor::new( + Box::new(DeltaOffsetCompressor::new(child.selection.compressor)), + encoding, + child_has_payload, + ); + candidates.push(OffsetFamilyCandidate { + estimated_wire_bytes: offset_wire_bytes( + &selection.encoding, + selection.has_payload, + &estimated_payload_bytes, + cost, + ), + selection, + estimated_payload_bytes, + transform_depth: 1, + decode_cpu_rank: 2, + stable_rank: 5, + }); + } + + candidates + .into_iter() + .reduce(|best, candidate| { + if candidate.is_better_than(&best) { + candidate + } else { + best + } + }) + .map(|candidate| candidate.finish(members)) + .ok_or_else(|| { + Error::internal("No applicable block family compression candidate".to_string()) + }) +} + +fn offset_leaf_candidate( + stats: &[&SequenceStats], + field_params: &CompressionFieldParams, + cost: BlockCost, +) -> Result { + let value_type = stats[0].value_type; + if stats.iter().all(|stats| stats.len == 0) { + return Ok(offset_metadata_candidate( + constant_selection(value_type, None), + stats.len(), + cost, + 0, + )); + } + if let Some(value) = common_constant(stats) { + return Ok(offset_metadata_candidate( + constant_selection(value_type, Some(value)), + stats.len(), + cost, + 0, + )); + } + if let Some((start, step)) = common_range(stats) { + return Ok(offset_metadata_candidate( + range_selection(value_type, start, step), + stats.len(), + cost, + 1, + )); + } + offset_direct_candidates(stats, field_params, cost, 0, false)? + .into_iter() + .reduce(|best, candidate| { + if candidate.is_better_than(&best) { + candidate + } else { + best + } + }) + .ok_or_else(|| Error::internal("No applicable block family leaf candidate".to_string())) +} + +fn offset_rle_candidate( + analysis: &OffsetFamilyAnalysis, + field_params: &CompressionFieldParams, + cost: BlockCost, +) -> Result> { + let Some(runs) = analysis.runs.as_ref() else { + return Ok(None); + }; + + let run_value_stats = runs.iter().map(|member| &member.values).collect::>(); + let raw_run_length_stats = runs + .iter() + .map(|member| &member.lengths) + .collect::>(); + let max_run_length = raw_run_length_stats + .iter() + .map(|stats| stats.max) + .max() + .unwrap_or(0); + let run_lengths_are_range = common_range(&raw_run_length_stats).is_some(); + let run_length_type = if run_lengths_are_range && max_run_length <= u32::MAX as u64 { + BlockValueType::UInt32 + } else if max_run_length <= u8::MAX as u64 { + BlockValueType::UInt8 + } else if max_run_length <= u16::MAX as u64 { + BlockValueType::UInt16 + } else if max_run_length <= u32::MAX as u64 { + BlockValueType::UInt32 + } else { + return Ok(None); + }; + let run_length_stats = raw_run_length_stats + .iter() + .map(|stats| SequenceStats { + value_type: run_length_type, + len: stats.len, + first: stats.first, + min: stats.min, + max: stats.max, + is_non_decreasing: stats.is_non_decreasing, + arithmetic_step: stats.arithmetic_step, + run_count: stats.run_count, + sample: None, + }) + .collect::>(); + let run_length_refs = run_length_stats.iter().collect::>(); + + let values = offset_leaf_candidate(&run_value_stats, field_params, BlockCost::default())?; + let mut lengths = offset_leaf_candidate(&run_length_refs, field_params, BlockCost::default())?; + let lengths_are_metadata = matches!( + lengths.selection.encoding.compression.as_ref(), + Some(crate::format::pb21::compressive_encoding::Compression::Constant(_)) + | Some(crate::format::pb21::compressive_encoding::Compression::Range(_)) + ); + let values_are_flat = matches!( + values.selection.encoding.compression.as_ref(), + Some(crate::format::pb21::compressive_encoding::Compression::Flat(_)) + ); + let lengths_are_flat = matches!( + lengths.selection.encoding.compression.as_ref(), + Some(crate::format::pb21::compressive_encoding::Compression::Flat(_)) + ); + if !lengths_are_metadata && !values_are_flat && !lengths_are_flat { + lengths = offset_flat_candidate(&run_length_refs, BlockCost::default())?; + } + let has_payload = values.selection.has_payload || lengths.selection.has_payload; + let estimated_payload_bytes = values + .estimated_payload_bytes + .iter() + .zip(&lengths.estimated_payload_bytes) + .map(|(values, lengths)| { + if has_payload { + rle::BLOCK_FRAME_BYTES + .saturating_add(*values) + .saturating_add(*lengths) + } else { + 0 + } + }) + .collect::>(); + let encoding = ProtobufUtils21::rle( + values.selection.encoding.clone(), + lengths.selection.encoding.clone(), + ); + let selection = SelectedBlockCompressor::new( + Box::new(BlockRleCompressor::new( + analysis.members[0].values.value_type, + run_length_type, + values.selection.compressor, + lengths.selection.compressor, + )), + encoding, + has_payload, + ); + Ok(Some(OffsetFamilyCandidate { + estimated_wire_bytes: offset_wire_bytes( + &selection.encoding, + has_payload, + &estimated_payload_bytes, + cost, + ), + selection, + estimated_payload_bytes, + transform_depth: 1, + decode_cpu_rank: 3, + stable_rank: 3, + })) +} + +fn offset_dictionary_candidate( + analysis: &OffsetFamilyAnalysis, + field_params: &CompressionFieldParams, + cost: BlockCost, +) -> Result> { + if analysis.members[0].values.value_type != BlockValueType::UInt64 { + return Ok(None); + } + let Some(dictionary_items) = analysis.dictionary_items.as_ref() else { + return Ok(None); + }; + let total_values = analysis.members.iter().try_fold(0_u64, |total, member| { + total + .checked_add(member.values.len) + .ok_or_else(|| Error::invalid_input("Offset family cardinality overflows u64")) + })?; + if dictionary_items.len() < 2 + || dictionary_items.len() > MAX_DICTIONARY_ITEMS + || dictionary_items.len() as u64 >= total_values.div_ceil(2) + { + return Ok(None); + } + + let mut items_builder = SequenceStatsBuilder::without_sample(BlockValueType::UInt64); + for value in dictionary_items.iter().copied() { + items_builder.push(value); + } + let items_stats = items_builder.finish(); + let item_stats = (0..analysis.members.len()) + .map(|_| &items_stats) + .collect::>(); + let index_stats = analysis + .members + .iter() + .map(|member| SequenceStats { + value_type: BlockValueType::UInt32, + len: member.values.len, + first: Some(0), + min: 0, + max: dictionary_items.len().saturating_sub(1) as u64, + is_non_decreasing: false, + arithmetic_step: None, + run_count: member.values.len, + sample: None, + }) + .collect::>(); + let index_refs = index_stats.iter().collect::>(); + let indices = offset_leaf_candidate(&index_refs, field_params, BlockCost::default())?; + let items = offset_leaf_candidate(&item_stats, field_params, BlockCost::default())?; + let has_payload = indices.selection.has_payload || items.selection.has_payload; + let estimated_payload_bytes = indices + .estimated_payload_bytes + .iter() + .zip(&items.estimated_payload_bytes) + .map(|(indices, items)| { + if has_payload { + dictionary::BLOCK_FRAME_BYTES + .saturating_add(*indices) + .saturating_add(*items) + } else { + 0 + } + }) + .collect::>(); + let num_dictionary_items = u32::try_from(dictionary_items.len()) + .map_err(|_| Error::invalid_input("Dictionary item count exceeds u32::MAX"))?; + let encoding = ProtobufUtils21::dictionary( + indices.selection.encoding.clone(), + items.selection.encoding.clone(), + num_dictionary_items, + ); + let selection = SelectedBlockCompressor::new( + Box::new(DictionaryBlockCompressor::new( + BlockValueType::UInt64, + dictionary_items.clone(), + indices.selection.compressor, + items.selection.compressor, + )), + encoding, + has_payload, + ); + Ok(Some(OffsetFamilyCandidate { + estimated_wire_bytes: offset_wire_bytes( + &selection.encoding, + has_payload, + &estimated_payload_bytes, + cost, + ), + selection, + estimated_payload_bytes, + transform_depth: 1, + decode_cpu_rank: 3, + stable_rank: 4, + })) +} + +fn offset_direct_candidates( + stats: &[&SequenceStats], + field_params: &CompressionFieldParams, + cost: BlockCost, + stable_rank_base: u8, + allow_general: bool, +) -> Result> { + let combined = combine_sequence_stats(stats)?; + direct_candidates(&combined, field_params, stable_rank_base, allow_general)? + .into_iter() + .map(|candidate: Candidate| { + let estimator = candidate.estimator.ok_or_else(|| { + Error::internal( + "Direct block candidate is missing its payload estimator".to_string(), + ) + })?; + let estimated_payload_bytes = stats + .iter() + .map(|stats| estimate_payload(estimator, stats)) + .collect::>>()?; + let estimated_wire_bytes = offset_wire_bytes( + &candidate.selection.encoding, + candidate.selection.has_payload, + &estimated_payload_bytes, + cost, + ); + Ok(OffsetFamilyCandidate { + selection: candidate.selection, + estimated_payload_bytes, + estimated_wire_bytes, + transform_depth: candidate.transform_depth, + decode_cpu_rank: candidate.decode_cpu_rank, + stable_rank: candidate.stable_rank, + }) + }) + .collect() +} + +fn offset_flat_candidate( + stats: &[&SequenceStats], + cost: BlockCost, +) -> Result { + let value_type = stats[0].value_type; + if stats.iter().any(|stats| stats.value_type != value_type) { + return Err(Error::invalid_input( + "Cannot build a Flat family candidate with different value widths", + )); + } + let estimated_payload_bytes = stats + .iter() + .map(|stats| stats.raw_bytes()) + .collect::>(); + let selection = flat_selection(value_type); + Ok(OffsetFamilyCandidate { + estimated_wire_bytes: offset_wire_bytes( + &selection.encoding, + true, + &estimated_payload_bytes, + cost, + ), + selection, + estimated_payload_bytes, + transform_depth: 0, + decode_cpu_rank: 0, + stable_rank: 0, + }) +} + +fn offset_metadata_candidate( + selection: SelectedBlockCompressor, + num_blocks: usize, + cost: BlockCost, + stable_rank: u8, +) -> OffsetFamilyCandidate { + let estimated_payload_bytes = vec![0; num_blocks]; + OffsetFamilyCandidate { + estimated_wire_bytes: offset_wire_bytes( + &selection.encoding, + false, + &estimated_payload_bytes, + cost, + ), + selection, + estimated_payload_bytes, + transform_depth: 0, + decode_cpu_rank: 0, + stable_rank, + } +} + +fn offset_wire_bytes( + encoding: &CompressiveEncoding, + has_payload: bool, + payload_bytes: &[u64], + cost: BlockCost, +) -> u64 { + let payload_cost = payload_bytes.iter().fold(0_u64, |total, payload_bytes| { + total.saturating_add(cost.payload_wire_bytes(has_payload, *payload_bytes)) + }); + (encoding.encoded_len() as u64).saturating_add(payload_cost) +} + +pub(super) fn common_first(stats: &[&SequenceStats]) -> Option { + let first = stats.first()?.first?; + stats + .iter() + .all(|stats| stats.first == Some(first)) + .then_some(first) +} + +pub(super) fn common_constant(stats: &[&SequenceStats]) -> Option { + let value = stats.first()?.first?; + stats + .iter() + .all(|stats| stats.len > 0 && stats.min == value && stats.max == value) + .then_some(value) +} + +pub(super) fn common_range(stats: &[&SequenceStats]) -> Option<(u64, u64)> { + let first = stats.first()?; + let start = first.first?; + let step = first.arithmetic_step?; + if step == 0 { + return None; + } + stats + .iter() + .all(|stats| { + stats.len >= 2 && stats.first == Some(start) && stats.arithmetic_step == Some(step) + }) + .then_some((start, step)) +} + +pub(super) fn combine_sequence_stats(stats: &[&SequenceStats]) -> Result { + let value_type = stats[0].value_type; + let mut len = 0_u64; + let mut run_count = 0_u64; + let mut min = u64::MAX; + let mut max = 0_u64; + let mut sample = Vec::new(); + for stats in stats { + if stats.value_type != value_type { + return Err(Error::invalid_input( + "Cannot combine sequence stats with different value widths", + )); + } + len = len + .checked_add(stats.len) + .ok_or_else(|| Error::invalid_input("Combined sequence length overflows u64"))?; + run_count = run_count + .checked_add(stats.run_count) + .ok_or_else(|| Error::invalid_input("Combined run count overflows u64"))?; + if stats.len > 0 { + min = min.min(stats.min); + max = max.max(stats.max); + } + let remaining = GENERAL_SAMPLE_BYTES.saturating_sub(sample.len()); + let stats_sample = stats.sample_bytes(); + sample.extend_from_slice(&stats_sample[..stats_sample.len().min(remaining)]); + } + Ok(SequenceStats { + value_type, + len, + first: stats.first().and_then(|stats| stats.first), + min: if len == 0 { 0 } else { min }, + max, + is_non_decreasing: false, + arithmetic_step: None, + run_count, + sample: (!sample.is_empty()).then(|| Arc::from(sample)), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn fixed_u32(values: &[u32]) -> FixedWidthDataBlock { + FixedWidthDataBlock { + data: crate::buffer::LanceBuffer::reinterpret_vec(values.to_vec()), + bits_per_value: 32, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + } + } + + fn fixed_u64(values: &[u64]) -> FixedWidthDataBlock { + FixedWidthDataBlock { + data: crate::buffer::LanceBuffer::reinterpret_vec(values.to_vec()), + bits_per_value: 64, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + } + } + + #[test] + fn delta_range_handles_a_short_final_chunk() { + let lengths = (0..4_112_u32) + .map(|index| 4 + index % 64) + .collect::>(); + let mut offsets = Vec::with_capacity(lengths.len() + 1); + offsets.push(0_u32); + for length in lengths { + offsets.push(offsets.last().copied().unwrap() + length); + } + let ranges = (0..offsets.len() - 1) + .step_by(64) + .map(|start| start..(start + 65).min(offsets.len())) + .collect::>(); + let analysis = analyze_offset_members(&fixed_u32(&offsets), &ranges, false).unwrap(); + let delta_stats = analysis + .members + .iter() + .map(|analysis| analysis.deltas.as_ref().unwrap()) + .collect::>(); + assert_eq!(common_range(&delta_stats), Some((4, 1))); + } + + #[test] + fn selects_rle_without_materializing_losing_dictionary() { + use crate::encodings::physical::dictionary::BLOCK_MATERIALIZATION_COUNT; + use crate::format::pb21::compressive_encoding::Compression; + + let member = (0..32_u64) + .flat_map(|value| std::iter::repeat_n(value, 8)) + .collect::>(); + let mut values = member.clone(); + values.extend_from_slice(&member); + let ranges = vec![0..member.len(), member.len()..values.len()]; + + BLOCK_MATERIALIZATION_COUNT.with(|count| count.set(0)); + let family = select_offset_family( + fixed_u64(&values), + ranges, + &CompressionFieldParams::default(), + BlockCost::default(), + ) + .unwrap(); + assert!(matches!( + family.encoding().compression.as_ref(), + Some(Compression::Rle(_)) + )); + assert!(!family.has_payload()); + BLOCK_MATERIALIZATION_COUNT.with(|count| assert_eq!(count.get(), 0)); + + let payloads = family.compress_members().unwrap(); + assert!(payloads.iter().all(Option::is_none)); + BLOCK_MATERIALIZATION_COUNT.with(|count| assert_eq!(count.get(), 0)); + } +} diff --git a/rust/lance-encoding/src/encodings/physical/binary/offsets/selector.rs b/rust/lance-encoding/src/encodings/physical/binary/offsets/selector.rs new file mode 100644 index 00000000000..5c78feb8609 --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/binary/offsets/selector.rs @@ -0,0 +1,233 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Private bounded selector for concrete unsigned block compressors. + +use super::{SelectedBlockCompressor, statistics::*}; +use crate::{ + compression::{BlockValueType, block::encode_scalar}, + compression_config::CompressionFieldParams, + encodings::physical::{ + block::{CompressedBufferEncoder, CompressionConfig, CompressionScheme}, + constant::ConstantBlockCompressor, + general::{self, GeneralBlockCompressor}, + range::RangeEncoder, + value::FixedWidthBlockCompressor, + }, + format::ProtobufUtils21, +}; +use lance_core::{Error, Result}; + +#[cfg(feature = "bitpacking")] +use crate::compression::{ + BlockCompressor, + block::{BITPACK_CHUNK_VALUES, out_of_line_payload_bytes}, +}; +#[cfg(feature = "bitpacking")] +use crate::encodings::physical::bitpacking::{InlineBitpacking, OutOfLineBitpacking}; + +/// How a direct leaf candidate's payload is estimated for another sequence. +/// +/// This cost metadata exists only while comparing concrete codecs. +#[derive(Debug, Clone, Copy)] +pub(super) enum PayloadEstimator { + Flat, + #[cfg(feature = "bitpacking")] + InlineBitpacking, + #[cfg(feature = "bitpacking")] + OutOfLineBitpacking { + compressed_bits: u64, + }, + General { + config: CompressionConfig, + }, +} + +#[derive(Debug)] +pub(super) struct Candidate { + pub(super) selection: SelectedBlockCompressor, + pub(super) estimator: Option, + pub(super) transform_depth: u8, + pub(super) decode_cpu_rank: u8, + pub(super) stable_rank: u8, +} + +fn estimate_general_with_config(stats: &SequenceStats, config: CompressionConfig) -> Result { + let sample = stats.sample_bytes(); + if sample.is_empty() && stats.raw_bytes() != 0 { + return Err(Error::invalid_input( + "Cannot estimate General block compression without a sample", + )); + } + let compressed = general::compress_block(config, sample)?; + Ok( + if sample.is_empty() || sample.len() as u64 == stats.raw_bytes() { + compressed.len() as u64 + } else { + (compressed.len() as u64) + .saturating_mul(stats.raw_bytes()) + .div_ceil(sample.len() as u64) + }, + ) +} + +pub(super) fn direct_candidates( + stats: &SequenceStats, + field_params: &CompressionFieldParams, + stable_rank_base: u8, + allow_general: bool, +) -> Result> { + let mut candidates = vec![Candidate { + selection: flat_selection(stats.value_type), + estimator: Some(PayloadEstimator::Flat), + transform_depth: 0, + decode_cpu_rank: 0, + stable_rank: stable_rank_base, + }]; + + #[cfg(feature = "bitpacking")] + if stats.len > 0 && stats.required_bits() < stats.value_type.bits_per_value() { + let compressed_bits = stats.required_bits(); + let (compressor, encoding, estimator) = if stats.len <= BITPACK_CHUNK_VALUES { + ( + Box::new(InlineBitpacking::new(stats.value_type.bits_per_value())) + as Box, + ProtobufUtils21::inline_bitpacking(stats.value_type.bits_per_value(), None), + PayloadEstimator::InlineBitpacking, + ) + } else { + ( + Box::new(OutOfLineBitpacking::new( + compressed_bits, + stats.value_type.bits_per_value(), + )) as Box, + ProtobufUtils21::out_of_line_bitpacking( + stats.value_type.bits_per_value(), + ProtobufUtils21::flat(compressed_bits, None), + ), + PayloadEstimator::OutOfLineBitpacking { compressed_bits }, + ) + }; + candidates.push(Candidate { + selection: SelectedBlockCompressor::new(compressor, encoding, true), + estimator: Some(estimator), + transform_depth: 1, + decode_cpu_rank: 1, + stable_rank: stable_rank_base.saturating_add(1), + }); + } + + if allow_general + && let Some((config, estimated_payload)) = estimate_general_payload(stats, field_params)? + && estimated_payload < stats.raw_bytes() + { + let child_encoding = ProtobufUtils21::flat(stats.value_type.bits_per_value(), None); + let encoding = ProtobufUtils21::wrapped(config, child_encoding)?; + let child = Box::new(FixedWidthBlockCompressor::new(stats.value_type)); + let compressor = Box::new(GeneralBlockCompressor::new(child, config)); + candidates.push(Candidate { + selection: SelectedBlockCompressor::new(compressor, encoding, true), + estimator: Some(PayloadEstimator::General { config }), + transform_depth: 1, + decode_cpu_rank: 4, + stable_rank: stable_rank_base.saturating_add(2), + }); + } + Ok(candidates) +} + +pub(super) fn flat_selection(value_type: BlockValueType) -> SelectedBlockCompressor { + SelectedBlockCompressor::new( + Box::new(FixedWidthBlockCompressor::new(value_type)), + ProtobufUtils21::flat(value_type.bits_per_value(), None), + true, + ) +} + +pub(super) fn constant_selection( + value_type: BlockValueType, + value: Option, +) -> SelectedBlockCompressor { + let encoded_value = value.map(|value| { + encode_scalar(value, value_type) + .expect("selector value was observed in the declared value type") + }); + SelectedBlockCompressor::new( + Box::new(ConstantBlockCompressor::new(value_type, value)), + ProtobufUtils21::constant(encoded_value), + false, + ) +} + +pub(super) fn range_selection( + value_type: BlockValueType, + start: u64, + step: u64, +) -> SelectedBlockCompressor { + SelectedBlockCompressor::new( + Box::new(RangeEncoder::new(value_type.bits_per_value(), start, step)), + ProtobufUtils21::range(value_type.bits_per_value(), start, step), + false, + ) +} + +fn explicit_general_config( + field_params: &CompressionFieldParams, +) -> Result> { + let Some(raw) = field_params.compression.as_deref() else { + return Ok(None); + }; + if raw == "none" { + return Ok(None); + } + let scheme: CompressionScheme = raw.parse()?; + match scheme { + CompressionScheme::Lz4 | CompressionScheme::Zstd => Ok(Some(CompressionConfig::new( + scheme, + field_params.compression_level, + ))), + CompressionScheme::None | CompressionScheme::Fsst => Err(Error::invalid_input(format!( + "Compression scheme '{raw}' is not supported for fixed-width block sequences" + ))), + } +} + +fn estimate_general_payload( + stats: &SequenceStats, + field_params: &CompressionFieldParams, +) -> Result> { + let config = match field_params.compression.as_deref() { + Some("none") => return Ok(None), + Some(_) => explicit_general_config(field_params)?, + None if stats.raw_bytes() > 32 * 1024 => { + Some(CompressedBufferEncoder::default().compressor.config()) + } + None => None, + }; + let Some(config) = config else { + return Ok(None); + }; + let sample = stats.sample_bytes(); + if sample.is_empty() { + return Ok(None); + } + Ok(Some((config, estimate_general_with_config(stats, config)?))) +} + +pub(super) fn estimate_payload(estimator: PayloadEstimator, stats: &SequenceStats) -> Result { + match estimator { + PayloadEstimator::Flat => Ok(stats.raw_bytes()), + #[cfg(feature = "bitpacking")] + PayloadEstimator::InlineBitpacking => Ok((1 + + (BITPACK_CHUNK_VALUES * stats.required_bits()) / stats.value_type.bits_per_value()) + .saturating_mul(stats.value_type.bytes_per_value() as u64)), + #[cfg(feature = "bitpacking")] + PayloadEstimator::OutOfLineBitpacking { compressed_bits } => { + if stats.required_bits() > compressed_bits { + return Ok(u64::MAX); + } + out_of_line_payload_bytes(stats.value_type, stats.len, compressed_bits) + } + PayloadEstimator::General { config } => estimate_general_with_config(stats, config), + } +} diff --git a/rust/lance-encoding/src/encodings/physical/binary/offsets/statistics.rs b/rust/lance-encoding/src/encodings/physical/binary/offsets/statistics.rs new file mode 100644 index 00000000000..01262264d97 --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/binary/offsets/statistics.rs @@ -0,0 +1,178 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use std::sync::Arc; + +use crate::compression::{BlockValueType, block::MAX_DICTIONARY_ITEMS}; + +#[derive(Debug, Clone)] +pub struct SequenceStats { + pub value_type: BlockValueType, + pub len: u64, + pub first: Option, + pub min: u64, + pub max: u64, + pub is_non_decreasing: bool, + pub arithmetic_step: Option, + pub run_count: u64, + pub sample: Option>, +} + +impl SequenceStats { + pub fn raw_bytes(&self) -> u64 { + self.len + .saturating_mul(self.value_type.bytes_per_value() as u64) + } + + #[cfg(feature = "bitpacking")] + pub fn required_bits(&self) -> u64 { + if self.max == 0 { + 1 + } else { + u64::from(u64::BITS - self.max.leading_zeros()) + } + } + + pub fn sample_bytes(&self) -> &[u8] { + self.sample.as_deref().unwrap_or_default() + } +} + +pub struct SequenceStatsBuilder { + value_type: BlockValueType, + len: u64, + first: Option, + previous: Option, + min: u64, + max: u64, + is_non_decreasing: bool, + step: Option, + is_arithmetic: bool, + run_count: u64, + sample: Vec, + sample_limit: usize, +} + +impl SequenceStatsBuilder { + pub fn without_sample(value_type: BlockValueType) -> Self { + Self::with_sample_limit(value_type, 0) + } + + pub fn with_sample_limit(value_type: BlockValueType, sample_limit: usize) -> Self { + Self { + value_type, + len: 0, + first: None, + previous: None, + min: u64::MAX, + max: 0, + is_non_decreasing: true, + step: None, + is_arithmetic: true, + run_count: 0, + sample: Vec::new(), + sample_limit, + } + } + + #[inline(always)] + pub fn push(&mut self, value: u64) { + if self.first.is_none() { + self.first = Some(value); + self.run_count = 1; + } + if let Some(previous) = self.previous { + self.is_non_decreasing &= value >= previous; + if value != previous { + self.run_count += 1; + } + let difference = value.checked_sub(previous); + match (self.step, difference) { + (None, Some(difference)) if self.len == 1 => self.step = Some(difference), + (Some(step), Some(difference)) if step == difference => {} + _ => self.is_arithmetic = false, + } + } + self.previous = Some(value); + self.min = self.min.min(value); + self.max = self.max.max(value); + self.len += 1; + + if self.sample.len() < self.sample_limit { + let bytes = value.to_le_bytes(); + let width = self.value_type.bytes_per_value(); + let remaining = self.sample_limit - self.sample.len(); + self.sample + .extend_from_slice(&bytes[..width.min(remaining)]); + } + } + + pub fn finish(self) -> SequenceStats { + let arithmetic_step = if self.len >= 2 && self.is_arithmetic { + self.step + } else { + None + }; + SequenceStats { + value_type: self.value_type, + len: self.len, + first: self.first, + min: if self.len == 0 { 0 } else { self.min }, + max: self.max, + is_non_decreasing: self.is_non_decreasing, + arithmetic_step, + run_count: self.run_count, + sample: (!self.sample.is_empty()).then(|| Arc::from(self.sample)), + } + } +} + +pub struct BoundedDistinctStatsBuilder { + items: Option>, +} + +impl BoundedDistinctStatsBuilder { + pub fn new() -> Self { + Self { + items: Some(Vec::new()), + } + } + + #[inline(always)] + pub fn push(&mut self, value: u64) { + let Some(items) = self.items.as_mut() else { + return; + }; + match items.binary_search(&value) { + Ok(_) => {} + Err(_) if items.len() == MAX_DICTIONARY_ITEMS => self.items = None, + Err(position) => { + items.insert(position, value); + } + } + } + + pub fn finish(self) -> Option> { + self.items.map(Arc::from) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn distinct_collection_stops_at_the_dictionary_bound() { + let mut bounded = BoundedDistinctStatsBuilder::new(); + for value in 0..MAX_DICTIONARY_ITEMS as u64 { + bounded.push(value); + } + assert_eq!(bounded.finish().unwrap().len(), MAX_DICTIONARY_ITEMS); + + let mut oversized = BoundedDistinctStatsBuilder::new(); + for value in 0..=MAX_DICTIONARY_ITEMS as u64 { + oversized.push(value); + } + assert!(oversized.finish().is_none()); + } +} diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index be0b747e7dc..9f049758c80 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -16,7 +16,7 @@ //! we can easily jump to the correct value. use arrow_array::types::UInt64Type; -use arrow_array::{Array, PrimitiveArray}; +use arrow_array::{Array, PrimitiveArray, UInt64Array}; use arrow_buffer::ArrowNativeType; use byteorder::{ByteOrder, LittleEndian}; use lance_bitpacking::BitPacking; @@ -24,20 +24,33 @@ use lance_bitpacking::BitPacking; use lance_core::{Error, Result}; use crate::buffer::LanceBuffer; -use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor}; +use crate::compression::{ + BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, + block::{ + validate_inline_bitpacking_payload, validate_out_of_line_payload, visit_unsigned_values, + }, + require_block_payload, +}; use crate::data::BlockInfo; use crate::data::{DataBlock, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::pb21::CompressiveEncoding; use crate::format::{ProtobufUtils21, pb21}; use crate::statistics::{GetStat, Stat}; use bytemuck::{AnyBitPattern, cast_slice}; +#[cfg(test)] +use std::cell::Cell; const LOG_ELEMS_PER_CHUNK: u8 = 10; const ELEMS_PER_CHUNK: u64 = 1 << LOG_ELEMS_PER_CHUNK; +#[cfg(test)] +thread_local! { + pub(crate) static BLOCK_WIDTH_VALIDATION_COUNT: Cell = const { Cell::new(0) }; +} + #[derive(Debug, Default)] pub struct InlineBitpacking { uncompressed_bit_width: u64, @@ -215,7 +228,11 @@ impl InlineBitpacking { } impl MiniBlockCompressor for InlineBitpacking { - fn compress(&self, chunk: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + chunk: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => Ok(self.chunk_data(fixed_width)), _ => Err(Error::invalid_input_source( @@ -230,10 +247,42 @@ impl MiniBlockCompressor for InlineBitpacking { } impl BlockCompressor for InlineBitpacking { - fn compress(&self, data: DataBlock) -> Result { - let fixed_width = data.as_fixed_width().unwrap(); + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(fixed_width) = data else { + return Err(Error::invalid_input( + "Inline bitpacking requires fixed-width data", + )); + }; + if fixed_width.bits_per_value != self.uncompressed_bit_width { + return Err(Error::invalid_input(format!( + "Inline bitpacking expects {}-bit values, got {}", + self.uncompressed_bit_width, fixed_width.bits_per_value + ))); + } + let value_type = BlockValueType::from_bits(self.uncompressed_bit_width)?; + if fixed_width.num_values == 0 || fixed_width.num_values > ELEMS_PER_CHUNK { + return Err(Error::invalid_input(format!( + "Inline block bitpacking requires 1..={ELEMS_PER_CHUNK} values, got {}", + fixed_width.num_values + ))); + } + let mut max_value = 0_u64; + visit_unsigned_values(&fixed_width, value_type, |value| { + max_value = max_value.max(value); + Ok(()) + })?; + let compressed_bit_width = u64::from(u64::BITS - max_value.leading_zeros()).max(1); + fixed_width.block_info.0.write().unwrap().insert( + Stat::BitWidth, + std::sync::Arc::new(UInt64Array::from(vec![compressed_bit_width])), + ); let (chunked, _) = self.chunk_data(fixed_width); - Ok(chunked.data.into_iter().next().unwrap()) + chunked + .data + .into_iter() + .next() + .map(Some) + .ok_or_else(|| Error::internal("Inline bitpacking produced no payload".to_string())) } } @@ -261,7 +310,13 @@ impl MiniBlockDecompressor for InlineBitpacking { } impl BlockDecompressor for InlineBitpacking { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Inline bitpacking")?; + validate_inline_bitpacking_payload( + &data, + BlockValueType::from_bits(self.uncompressed_bit_width)?, + num_values, + )?; match self.uncompressed_bit_width { 8 => Self::unchunk::(data, num_values), 16 => Self::unchunk::(data, num_values), @@ -473,27 +528,45 @@ impl OutOfLineBitpacking { } impl BlockCompressor for OutOfLineBitpacking { - fn compress(&self, data: DataBlock) -> Result { - let fixed_width = data.as_fixed_width().unwrap(); + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(fixed_width) = data else { + return Err(Error::invalid_input( + "Out-of-line bitpacking requires fixed-width data", + )); + }; + if fixed_width.bits_per_value != self.uncompressed_bit_width { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking expects {}-bit values, got {}", + self.uncompressed_bit_width, fixed_width.bits_per_value + ))); + } + validate_frozen_block_width(&fixed_width, self.compressed_bit_width)?; let compressed = match fixed_width.bits_per_value { 8 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), 16 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), 32 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), 64 => bitpack_out_of_line::(fixed_width, self.compressed_bit_width as usize), - _ => panic!("Bitpacking word size must be 8,16,32,64"), + _ => unreachable!("bitpacking word size was validated"), }; - Ok(compressed) + Ok(Some(compressed)) } } impl BlockDecompressor for OutOfLineBitpacking { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Out-of-line bitpacking")?; + validate_out_of_line_payload( + &data, + BlockValueType::from_bits(self.uncompressed_bit_width)?, + num_values, + self.compressed_bit_width, + )?; let word_size = match self.uncompressed_bit_width { 8 => std::mem::size_of::(), 16 => std::mem::size_of::(), 32 => std::mem::size_of::(), 64 => std::mem::size_of::(), - _ => panic!("Bitpacking word size must be 8,16,32,64"), + _ => unreachable!("bitpacking word size was validated"), }; debug_assert_eq!(data.len() % word_size, 0); let total_words = (data.len() / word_size) as u64; @@ -531,6 +604,38 @@ impl BlockDecompressor for OutOfLineBitpacking { } } +fn validate_frozen_block_width( + block: &FixedWidthDataBlock, + compressed_bits_per_value: u64, +) -> Result<()> { + #[cfg(test)] + BLOCK_WIDTH_VALIDATION_COUNT.with(|count| count.set(count.get().saturating_add(1))); + + let value_type = BlockValueType::from_bits(block.bits_per_value)?; + if compressed_bits_per_value >= value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking width {compressed_bits_per_value} must be less than the {}-bit input width", + value_type.bits_per_value() + ))); + } + let max_value = if compressed_bits_per_value == 0 { + 0 + } else { + (1_u64 << compressed_bits_per_value) - 1 + }; + let mut index = 0_u64; + visit_unsigned_values(block, value_type, |value| { + if value > max_value { + let required_bits = u64::from(u64::BITS - value.leading_zeros()); + return Err(Error::invalid_input(format!( + "Out-of-line bitpacking codec uses {compressed_bits_per_value} bits but input value {value} at index {index} requires {required_bits} bits" + ))); + } + index += 1; + Ok(()) + }) +} + #[cfg(test)] mod test { use std::{collections::HashMap, sync::Arc}; @@ -542,7 +647,7 @@ mod test { use super::{ELEMS_PER_CHUNK, InlineBitpacking, bitpack_out_of_line, unpack_out_of_line}; use crate::{ buffer::LanceBuffer, - compression::MiniBlockDecompressor, + compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor}, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, testing::{TestCases, check_round_trip_encoding_of_data}, version::LanceFileVersion, @@ -567,6 +672,27 @@ mod test { assert_eq!(block.data.len(), 0); } + #[test] + fn test_inline_block_bitpacking_computes_its_own_bit_width() { + let codec = InlineBitpacking::new(32); + let input = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_u32, 1, 7]), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::new(), + }); + + let payload = BlockCompressor::compress(&codec, input).unwrap().unwrap(); + let decoded = BlockDecompressor::decompress(&codec, Some(payload), 3).unwrap(); + let DataBlock::FixedWidth(decoded) = decoded else { + panic!("Expected FixedWidth block"); + }; + assert_eq!( + decoded.data.borrow_to_typed_view::().as_ref(), + &[0, 1, 7] + ); + } + #[test_log::test(tokio::test)] async fn test_miniblock_bitpack() { let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1); diff --git a/rust/lance-encoding/src/encodings/physical/block.rs b/rust/lance-encoding/src/encodings/physical/block.rs index a1f5bdb3fdd..17af022ae71 100644 --- a/rust/lance-encoding/src/encodings/physical/block.rs +++ b/rust/lance-encoding/src/encodings/physical/block.rs @@ -34,7 +34,7 @@ use crate::format::{ }; use crate::{ buffer::LanceBuffer, - compression::VariablePerValueDecompressor, + compression::{VariablePerValueDecompressor, require_block_payload}, data::{BlockInfo, DataBlock, VariableWidthBlock}, encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}, }; @@ -130,12 +130,40 @@ impl FromStr for CompressionScheme { pub trait BufferCompressor: std::fmt::Debug + Send + Sync { fn compress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()>; fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec) -> Result<()>; + /// Decompresses at most `max_output_len` bytes. + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()>; + /// Decompresses exactly `expected_output_len` bytes without trusting an + /// embedded size prefix to choose the allocation size. + fn decompress_exact( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + expected_output_len: usize, + ) -> Result<()> { + let start = output_buf.len(); + self.decompress_bounded(input_buf, output_buf, expected_output_len)?; + let actual_output_len = output_buf.len().checked_sub(start).ok_or_else(|| { + Error::internal("Buffer decompressor shortened its output".to_string()) + })?; + if actual_output_len != expected_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Decompression produced {actual_output_len} bytes, expected exactly {expected_output_len}" + ))); + } + Ok(()) + } fn config(&self) -> CompressionConfig; } #[cfg(feature = "zstd")] mod zstd { - use std::io::{Cursor, Write}; + use std::io::{Cursor, Read, Write}; use std::sync::{Mutex, OnceLock}; use super::*; @@ -279,6 +307,101 @@ mod zstd { Ok(()) } + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()> { + if self.is_raw_stream_format(input_buf) { + let read_limit = max_output_len.checked_add(1).ok_or_else(|| { + Error::invalid_input("Zstd maximum output length overflows usize") + })?; + let read_limit_u64 = u64::try_from(read_limit).map_err(|_| { + Error::invalid_input("Zstd maximum output length does not fit u64") + })?; + let start = output_buf.len(); + output_buf.try_reserve_exact(read_limit).map_err(|error| { + Error::invalid_input(format!( + "Zstd could not reserve {read_limit} bounded output bytes: {error}" + )) + })?; + let result = (|| { + let decoder = ::zstd::stream::read::Decoder::new(Cursor::new(input_buf))?; + let mut bounded = decoder.take(read_limit_u64); + bounded.read_to_end(output_buf) + })(); + let actual_output_len = match result { + Ok(actual_output_len) => actual_output_len, + Err(error) => { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd decompression failed: {error}" + ))); + } + }; + if actual_output_len > max_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd output exceeds the {max_output_len}-byte limit" + ))); + } + return Ok(()); + } + + const LENGTH_PREFIX_SIZE: usize = 8; + if input_buf.len() < LENGTH_PREFIX_SIZE { + return Err(Error::invalid_input(format!( + "Length-prefixed Zstd payload has {} bytes, shorter than its {LENGTH_PREFIX_SIZE}-byte prefix", + input_buf.len() + ))); + } + let declared_output_len = + u64::from_le_bytes(input_buf[..LENGTH_PREFIX_SIZE].try_into().map_err(|_| { + Error::invalid_input("Length-prefixed Zstd size prefix is truncated") + })?); + let declared_output_len = usize::try_from(declared_output_len).map_err(|_| { + Error::invalid_input("Length-prefixed Zstd output length does not fit usize") + })?; + if declared_output_len > max_output_len { + return Err(Error::invalid_input(format!( + "Length-prefixed Zstd payload declares {declared_output_len} output bytes, exceeding the {max_output_len}-byte limit" + ))); + } + + let start = output_buf.len(); + let end = start + .checked_add(declared_output_len) + .ok_or_else(|| Error::invalid_input("Zstd output buffer length overflows usize"))?; + output_buf + .try_reserve_exact(declared_output_len) + .map_err(|error| { + Error::invalid_input(format!( + "Zstd could not reserve {declared_output_len} output bytes: {error}" + )) + })?; + output_buf.resize(end, 0); + let decompressed_size = match decompress_to_buffer( + &input_buf[LENGTH_PREFIX_SIZE..], + &mut output_buf[start..end], + ) { + Ok(decompressed_size) => decompressed_size, + Err(error) => { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd decompression failed: {error}" + ))); + } + }; + if decompressed_size != declared_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "Zstd decompressed {decompressed_size} bytes, but its prefix declares {declared_output_len}" + ))); + } + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::Zstd, @@ -347,6 +470,62 @@ mod lz4 { Ok(()) } + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()> { + const LENGTH_PREFIX_SIZE: usize = 4; + if input_buf.len() < LENGTH_PREFIX_SIZE { + return Err(Error::invalid_input(format!( + "LZ4 payload has {} bytes, shorter than its {LENGTH_PREFIX_SIZE}-byte prefix", + input_buf.len() + ))); + } + let declared_output_len = u32::from_le_bytes( + input_buf[..LENGTH_PREFIX_SIZE] + .try_into() + .map_err(|_| Error::invalid_input("LZ4 size prefix is truncated"))?, + ) as usize; + if declared_output_len > max_output_len { + return Err(Error::invalid_input(format!( + "LZ4 payload declares {declared_output_len} output bytes, exceeding the {max_output_len}-byte limit" + ))); + } + + let start = output_buf.len(); + let end = start + .checked_add(declared_output_len) + .ok_or_else(|| Error::invalid_input("LZ4 output buffer length overflows usize"))?; + output_buf + .try_reserve_exact(declared_output_len) + .map_err(|error| { + Error::invalid_input(format!( + "LZ4 could not reserve {declared_output_len} output bytes: {error}" + )) + })?; + output_buf.resize(end, 0); + let decompressed_size = + match ::lz4::block::decompress_to_buffer(input_buf, None, &mut output_buf[start..]) + { + Ok(decompressed_size) => decompressed_size, + Err(error) => { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "LZ4 decompression failed: {error}" + ))); + } + }; + if decompressed_size != declared_output_len { + output_buf.truncate(start); + return Err(Error::invalid_input(format!( + "LZ4 decompressed {decompressed_size} bytes, but its prefix declares {declared_output_len}" + ))); + } + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::Lz4, @@ -370,6 +549,30 @@ impl BufferCompressor for NoopBufferCompressor { Ok(()) } + fn decompress_bounded( + &self, + input_buf: &[u8], + output_buf: &mut Vec, + max_output_len: usize, + ) -> Result<()> { + if input_buf.len() > max_output_len { + return Err(Error::invalid_input(format!( + "Uncompressed payload has {} bytes, exceeding the {max_output_len}-byte limit", + input_buf.len() + ))); + } + output_buf + .try_reserve_exact(input_buf.len()) + .map_err(|error| { + Error::invalid_input(format!( + "Uncompressed payload could not reserve {} output bytes: {error}", + input_buf.len() + )) + })?; + output_buf.extend_from_slice(input_buf); + Ok(()) + } + fn config(&self) -> CompressionConfig { CompressionConfig { scheme: CompressionScheme::None, @@ -439,11 +642,12 @@ impl GeneralBlockDecompressor { } impl BlockDecompressor for GeneralBlockDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "General block compression")?; let mut decompressed = Vec::new(); self.compressor.decompress(&data, &mut decompressed)?; self.inner - .decompress(LanceBuffer::from(decompressed), num_values) + .decompress(Some(LanceBuffer::from(decompressed)), num_values) } } @@ -603,13 +807,19 @@ impl VariablePerValueDecompressor for CompressedBufferEncoder { } impl BlockCompressor for CompressedBufferEncoder { - fn compress(&self, data: DataBlock) -> Result { + fn compress(&self, data: DataBlock) -> Result> { let encoded = match data { DataBlock::FixedWidth(fixed_width) => fixed_width.data, DataBlock::VariableWidth(variable_width) => { // Wrap VariableEncoder to handle the encoding let encoder = VariableEncoder::default(); BlockCompressor::compress(&encoder, DataBlock::VariableWidth(variable_width))? + .ok_or_else(|| { + Error::internal( + "VariableEncoder returned no payload for general compression" + .to_string(), + ) + })? } _ => { return Err(Error::invalid_input_source( @@ -620,18 +830,19 @@ impl BlockCompressor for CompressedBufferEncoder { let mut compressed = Vec::new(); self.compressor.compress(&encoded, &mut compressed)?; - Ok(LanceBuffer::from(compressed)) + Ok(Some(LanceBuffer::from(compressed))) } } impl BlockDecompressor for CompressedBufferEncoder { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Compressed variable block")?; let mut decompressed = Vec::new(); self.compressor.decompress(&data, &mut decompressed)?; // Delegate to BinaryBlockDecompressor which handles the inline metadata let inner_decoder = BinaryBlockDecompressor::default(); - inner_decoder.decompress(LanceBuffer::from(decompressed), num_values) + inner_decoder.decompress(Some(LanceBuffer::from(decompressed)), num_values) } } @@ -640,6 +851,7 @@ mod tests { use super::*; use std::str::FromStr; + #[cfg(feature = "zstd")] use crate::encodings::physical::block::zstd::ZstdBufferCompressor; #[test] @@ -679,6 +891,32 @@ mod tests { .decompress(&compressed_data, &mut decompressed_data) .unwrap(); assert_eq!(input_data, decompressed_data.as_slice()); + + let mut exact = Vec::new(); + compressor + .decompress_exact(&compressed_data, &mut exact, input_data.len()) + .unwrap(); + assert_eq!(input_data, exact.as_slice()); + } + + #[test] + fn test_length_prefixed_zstd_rejects_untrusted_output_size() { + let compressor = ZstdBufferCompressor::new(0); + let mut compressed_data = Vec::new(); + compressor + .compress(b"bounded", &mut compressed_data) + .unwrap(); + compressed_data[..8].copy_from_slice(&u64::MAX.to_le_bytes()); + + let mut output = vec![7, 9]; + let error = compressor + .decompress_exact(&compressed_data, &mut output, 7) + .unwrap_err(); + assert!( + error.to_string().contains("exceeding the 7-byte limit"), + "unexpected error: {error}" + ); + assert_eq!(output, [7, 9]); } #[test] @@ -745,6 +983,22 @@ mod tests { .decompress(&compressed_data, &mut decompressed_data) .unwrap(); assert_eq!(input_data, decompressed_data.as_slice()); + + let mut exact = Vec::new(); + compressor + .decompress_exact(&compressed_data, &mut exact, input_data.len()) + .unwrap(); + assert_eq!(input_data, exact.as_slice()); + + let mut bounded = vec![3, 5]; + let error = compressor + .decompress_exact(&compressed_data, &mut bounded, input_data.len() - 1) + .unwrap_err(); + assert!( + error.to_string().contains("exceeds the 12-byte limit"), + "unexpected error: {error}" + ); + assert_eq!(bounded, [3, 5]); } } @@ -782,6 +1036,32 @@ mod tests { .decompress(&compressed_data, &mut decompressed_data) .unwrap(); assert_eq!(input_data, decompressed_data.as_slice()); + + let mut exact = Vec::new(); + compressor + .decompress_exact(&compressed_data, &mut exact, input_data.len()) + .unwrap(); + assert_eq!(input_data, exact.as_slice()); + } + + #[test] + fn test_lz4_rejects_untrusted_output_size() { + let compressor = Lz4BufferCompressor::default(); + let mut compressed_data = Vec::new(); + compressor + .compress(b"bounded", &mut compressed_data) + .unwrap(); + compressed_data[..4].copy_from_slice(&u32::MAX.to_le_bytes()); + + let mut output = vec![7, 9]; + let error = compressor + .decompress_exact(&compressed_data, &mut output, 7) + .unwrap_err(); + assert!( + error.to_string().contains("exceeding the 7-byte limit"), + "unexpected error: {error}" + ); + assert_eq!(output, [7, 9]); } #[test_log::test(tokio::test)] diff --git a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs index c2b7aac9b9c..b861098d3e4 100644 --- a/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs +++ b/rust/lance-encoding/src/encodings/physical/byte_stream_split.rs @@ -62,7 +62,7 @@ use crate::compression::MiniBlockDecompressor; use crate::compression_config::BssMode; use crate::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ - MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressor, + MiniBlockChunk, MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::ProtobufUtils21; use crate::format::pb21::CompressiveEncoding; @@ -107,7 +107,11 @@ impl ByteStreamSplitEncoder { } impl MiniBlockCompressor for ByteStreamSplitEncoder { - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + page: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match page { DataBlock::FixedWidth(data) => { let num_values = data.num_values; @@ -345,7 +349,9 @@ mod tests { }); // Compress - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress let decompressed = decompressor @@ -391,7 +397,9 @@ mod tests { }); // Compress - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress let decompressed = decompressor @@ -424,7 +432,9 @@ mod tests { }); // Compress empty data - let (compressed, _encoding) = encoder.compress(data_block).unwrap(); + let (compressed, _encoding) = encoder + .compress(data_block, MiniBlockCompressionContext::new(0, true, true)) + .unwrap(); // Decompress empty data let decompressed = decompressor.decompress(compressed.data, 0).unwrap(); diff --git a/rust/lance-encoding/src/encodings/physical/constant.rs b/rust/lance-encoding/src/encodings/physical/constant.rs index c3fa16863f4..9084567e9ee 100644 --- a/rust/lance-encoding/src/encodings/physical/constant.rs +++ b/rust/lance-encoding/src/encodings/physical/constant.rs @@ -5,11 +5,15 @@ use crate::{ buffer::LanceBuffer, - compression::{BlockDecompressor, FixedPerValueDecompressor}, - data::{AllNullDataBlock, ConstantDataBlock, DataBlock, FixedWidthDataBlock}, + compression::{ + 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, }; -use lance_core::Result; +use lance_core::{Error, Result}; /// A decompressor for constant-encoded data #[derive(Debug)] @@ -24,7 +28,7 @@ impl ConstantDecompressor { } impl BlockDecompressor for ConstantDecompressor { - fn decompress(&self, _data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, _data: Option, num_values: u64) -> Result { if let Some(scalar) = self.scalar.clone() { Ok(DataBlock::Constant(ConstantDataBlock { data: scalar, @@ -55,3 +59,131 @@ impl FixedPerValueDecompressor for ConstantDecompressor { .unwrap_or(0) } } + +/// Metadata-only fixed-width constant (or typed empty) block compressor. +#[derive(Debug)] +pub(crate) struct ConstantBlockCompressor { + value_type: BlockValueType, + value: Option, +} + +impl ConstantBlockCompressor { + pub(crate) fn new(value_type: BlockValueType, value: Option) -> Self { + Self { value_type, value } + } +} + +impl BlockCompressor for ConstantBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Constant block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Constant block compressor expects {}-bit values, got {}", + self.value_type.bits_per_value(), + data.bits_per_value + ))); + } + validate_fixed_payload_len( + &data.data, + self.value_type, + data.num_values, + "Constant block input", + )?; + match self.value { + None => { + if data.num_values != 0 { + return Err(Error::invalid_input( + "Typed empty block compressor received non-empty data", + )); + } + } + Some(expected) => { + if data.num_values == 0 { + return Err(Error::invalid_input( + "Constant block compressor received an empty sequence", + )); + } + macro_rules! check_values { + ($ty:ty) => { + for (index, actual) in + data.data.borrow_to_typed_slice::<$ty>().iter().enumerate() + { + if u64::from(*actual) != expected { + return Err(Error::invalid_input(format!( + "Constant block expects {expected}, got {actual} at index {index}" + ))); + } + } + }; + } + match self.value_type { + BlockValueType::UInt8 => check_values!(u8), + BlockValueType::UInt16 => check_values!(u16), + BlockValueType::UInt32 => check_values!(u32), + BlockValueType::UInt64 => check_values!(u64), + } + } + } + Ok(None) + } +} + +/// Metadata-only fixed-width constant (or typed empty) block decompressor. +#[derive(Debug)] +pub(crate) struct ConstantBlockDecompressor { + value_type: BlockValueType, + value: Option, +} + +impl ConstantBlockDecompressor { + pub(crate) fn new(value_type: BlockValueType, value: Option) -> Self { + Self { value_type, value } + } +} + +impl BlockDecompressor for ConstantBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + require_no_block_payload(data, "Constant block")?; + let output = match self.value { + None => { + if num_values != 0 { + return Err(Error::invalid_input(format!( + "Typed empty block cannot represent {num_values} values" + ))); + } + LanceBuffer::empty() + } + Some(value) => { + if num_values == 0 { + return Err(Error::invalid_input( + "Non-empty Constant descriptor cannot represent an empty block", + )); + } + macro_rules! repeat { + ($ty:ty) => {{ + let mut values = + try_vec_with_capacity::<$ty>(num_values, "Constant block output")?; + values.resize(num_values as usize, value as $ty); + LanceBuffer::reinterpret_vec(values) + }}; + } + match self.value_type { + BlockValueType::UInt8 => repeat!(u8), + BlockValueType::UInt16 => repeat!(u16), + BlockValueType::UInt32 => repeat!(u32), + BlockValueType::UInt64 => repeat!(u64), + } + } + }; + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + data: output, + bits_per_value: self.value_type.bits_per_value(), + num_values, + block_info: BlockInfo::new(), + })) + } +} diff --git a/rust/lance-encoding/src/encodings/physical/delta.rs b/rust/lance-encoding/src/encodings/physical/delta.rs new file mode 100644 index 00000000000..5fade5f26b5 --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/delta.rs @@ -0,0 +1,276 @@ +// 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. +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/dictionary.rs b/rust/lance-encoding/src/encodings/physical/dictionary.rs new file mode 100644 index 00000000000..50225b55e6e --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/dictionary.rs @@ -0,0 +1,319 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Dictionary codec for bounded unsigned block sequences. + +#[cfg(test)] +use std::cell::Cell; +use std::{collections::BTreeMap, sync::Arc}; + +use crate::{ + buffer::LanceBuffer, + compression::{ + BlockCompressor, BlockDecompressor, BlockValueType, + block::{ + fixed_block, fixed_from_u64_values, read_unsigned_values, validate_fixed_payload_len, + visit_unsigned_values, + }, + }, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + encodings::physical::try_vec_with_capacity, +}; +use lance_core::{Error, Result}; + +pub(crate) const BLOCK_FRAME_BYTES: u64 = 16; + +#[cfg(test)] +thread_local! { + pub(crate) static BLOCK_MATERIALIZATION_COUNT: Cell = const { Cell::new(0) }; +} + +/// Dictionary compressor that owns its indices and items compressors. +#[derive(Debug)] +pub(crate) struct DictionaryBlockCompressor { + value_type: BlockValueType, + dictionary_items: Arc<[u64]>, + indices: Box, + items: Box, +} + +impl DictionaryBlockCompressor { + pub(crate) fn new( + value_type: BlockValueType, + dictionary_items: Arc<[u64]>, + indices: Box, + items: Box, + ) -> Self { + Self { + value_type, + dictionary_items, + indices, + items, + } + } +} + +impl BlockCompressor for DictionaryBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + #[cfg(test)] + BLOCK_MATERIALIZATION_COUNT.with(|count| count.set(count.get().saturating_add(1))); + + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Dictionary block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Dictionary compressor expects {}-bit values, got {}", + self.value_type.bits_per_value(), + data.bits_per_value + ))); + } + let dictionary_index = self + .dictionary_items + .iter() + .enumerate() + .map(|(index, value)| (*value, index as u32)) + .collect::>(); + let mut encoded_indices = + try_vec_with_capacity::(data.num_values, "Dictionary indices")?; + let mut position = 0_u64; + visit_unsigned_values(&data, self.value_type, |value| { + encoded_indices.push(*dictionary_index.get(&value).ok_or_else(|| { + Error::invalid_input(format!( + "Dictionary compressor does not contain input value {value} at position {position}" + )) + })?); + position += 1; + Ok(()) + })?; + + let indices_block = FixedWidthDataBlock { + bits_per_value: 32, + data: LanceBuffer::reinterpret_vec(encoded_indices), + num_values: data.num_values, + block_info: BlockInfo::default(), + }; + let items_block = + fixed_from_u64_values(&self.dictionary_items, self.value_type, "Dictionary items")?; + let indices_payload = self + .indices + .compress(DataBlock::FixedWidth(indices_block))?; + let items_payload = self.items.compress(DataBlock::FixedWidth(items_block))?; + if indices_payload.is_none() && items_payload.is_none() { + return Ok(None); + } + let indices_payload = indices_payload.unwrap_or_else(LanceBuffer::empty); + let items_payload = items_payload.unwrap_or_else(LanceBuffer::empty); + + let mut output = try_frame(indices_payload.len(), items_payload.len())?; + output.extend_from_slice(&(indices_payload.len() as u64).to_le_bytes()); + output.extend_from_slice(&(items_payload.len() as u64).to_le_bytes()); + output.extend_from_slice(&indices_payload); + output.extend_from_slice(&items_payload); + Ok(Some(LanceBuffer::from(output))) + } +} + +/// Dictionary decompressor that owns its indices and items decompressors. +#[derive(Debug)] +pub(crate) struct DictionaryBlockDecompressor { + value_type: BlockValueType, + num_dictionary_items: u32, + indices: Box, + items: Box, + indices_have_payload: bool, + items_have_payload: bool, +} + +impl DictionaryBlockDecompressor { + pub(crate) fn new( + value_type: BlockValueType, + num_dictionary_items: u32, + indices: Box, + items: Box, + indices_have_payload: bool, + items_have_payload: bool, + ) -> Self { + Self { + value_type, + num_dictionary_items, + indices, + items, + indices_have_payload, + items_have_payload, + } + } +} + +impl BlockDecompressor for DictionaryBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let has_payload = self.indices_have_payload || self.items_have_payload; + let (indices_payload, items_payload) = if let Some(data) = data { + if !has_payload { + return Err(Error::invalid_input( + "Metadata-only Dictionary expects no payload", + )); + } + if data.len() < BLOCK_FRAME_BYTES as usize { + return Err(Error::invalid_input(format!( + "Dictionary payload has {} bytes, shorter than its {BLOCK_FRAME_BYTES}-byte header", + data.len() + ))); + } + let indices_size = + u64::from_le_bytes(data[..8].try_into().expect("header length was checked")); + let items_size = + u64::from_le_bytes(data[8..16].try_into().expect("header length was checked")); + let indices_size = usize::try_from(indices_size).map_err(|_| { + Error::invalid_input("Dictionary indices payload length does not fit usize") + })?; + let items_size = usize::try_from(items_size).map_err(|_| { + Error::invalid_input("Dictionary items payload length does not fit usize") + })?; + let indices_start = BLOCK_FRAME_BYTES as usize; + let items_start = indices_start + .checked_add(indices_size) + .ok_or_else(|| Error::invalid_input("Dictionary indices payload end overflows"))?; + let end = items_start + .checked_add(items_size) + .ok_or_else(|| Error::invalid_input("Dictionary items payload end overflows"))?; + if end != data.len() { + return Err(Error::invalid_input(format!( + "Dictionary framing describes {end} bytes, payload has {}", + data.len() + ))); + } + if !self.indices_have_payload && indices_size != 0 { + return Err(Error::invalid_input(format!( + "Metadata-only Dictionary indices child has {indices_size} framed payload bytes" + ))); + } + if !self.items_have_payload && items_size != 0 { + return Err(Error::invalid_input(format!( + "Metadata-only Dictionary items child has {items_size} framed payload bytes" + ))); + } + ( + self.indices_have_payload + .then(|| data.slice_with_length(indices_start, indices_size)), + self.items_have_payload + .then(|| data.slice_with_length(items_start, items_size)), + ) + } else { + if has_payload { + return Err(Error::invalid_input("Dictionary requires one payload")); + } + (None, None) + }; + + let indices = self.indices.decompress(indices_payload, num_values)?; + let items = self + .items + .decompress(items_payload, u64::from(self.num_dictionary_items))?; + let DataBlock::FixedWidth(indices) = indices else { + return Err(Error::invalid_input( + "Dictionary indices decoded to a non fixed-width block", + )); + }; + let DataBlock::FixedWidth(items) = items else { + return Err(Error::invalid_input( + "Dictionary items decoded to a non fixed-width block", + )); + }; + validate_fixed_payload_len( + &indices.data, + BlockValueType::UInt32, + num_values, + "Dictionary indices", + )?; + validate_fixed_payload_len( + &items.data, + self.value_type, + u64::from(self.num_dictionary_items), + "Dictionary items", + )?; + let indices = indices.data.borrow_to_typed_slice::(); + let items = read_unsigned_values(&items, self.value_type)?; + let output = match self.value_type { + BlockValueType::UInt8 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value as u8, + )?; + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt16 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value as u16, + )?; + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt32 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value as u32, + )?; + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt64 => { + let mut output = try_vec_with_capacity::(num_values, "Dictionary output")?; + append_items( + &mut output, + &indices, + &items, + self.num_dictionary_items, + |value| value, + )?; + LanceBuffer::reinterpret_vec(output) + } + }; + Ok(fixed_block(self.value_type, num_values, output)) + } +} + +fn append_items( + output: &mut Vec, + indices: &[u32], + items: &[u64], + num_dictionary_items: u32, + convert: impl Fn(u64) -> T, +) -> Result<()> { + for (position, index) in indices.iter().enumerate() { + let value = *items.get(*index as usize).ok_or_else(|| { + Error::invalid_input(format!( + "Dictionary index {index} at position {position} is out of bounds for {num_dictionary_items} items" + )) + })?; + output.push(convert(value)); + } + Ok(()) +} + +fn try_frame(indices_payload_bytes: usize, items_payload_bytes: usize) -> Result> { + let capacity = (BLOCK_FRAME_BYTES as usize) + .checked_add(indices_payload_bytes) + .and_then(|capacity| capacity.checked_add(items_payload_bytes)) + .ok_or_else(|| Error::invalid_input("Dictionary frame length overflows usize"))?; + let mut output = Vec::new(); + output.try_reserve_exact(capacity).map_err(|error| { + Error::invalid_input(format!( + "Dictionary could not reserve {capacity} frame bytes: {error}" + )) + })?; + Ok(output) +} diff --git a/rust/lance-encoding/src/encodings/physical/fsst.rs b/rust/lance-encoding/src/encodings/physical/fsst.rs index 8c1fe4141df..295ffcd88b1 100644 --- a/rust/lance-encoding/src/encodings/physical/fsst.rs +++ b/rust/lance-encoding/src/encodings/physical/fsst.rs @@ -23,7 +23,7 @@ use crate::{ data::{BlockInfo, DataBlock, VariableWidthBlock}, encodings::logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, }, format::{ ProtobufUtils21, @@ -138,7 +138,11 @@ impl FsstMiniBlockEncoder { } impl MiniBlockCompressor for FsstMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { let compressed = FsstCompressed::fsst_compress(data)?; let data_block = DataBlock::VariableWidth(compressed.data); @@ -148,7 +152,7 @@ impl MiniBlockCompressor for FsstMiniBlockEncoder { as Box; let (binary_miniblock_compressed, binary_array_encoding) = - binary_compressor.compress(data_block)?; + binary_compressor.compress(data_block, context)?; Ok(( binary_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/general.rs b/rust/lance-encoding/src/encodings/physical/general.rs index 53c61928870..ff66363affd 100644 --- a/rust/lance-encoding/src/encodings/physical/general.rs +++ b/rust/lance-encoding/src/encodings/physical/general.rs @@ -6,14 +6,97 @@ use log::trace; use crate::{ Result, buffer::LanceBuffer, - compression::MiniBlockDecompressor, + compression::{ + BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, + require_block_payload, + }, data::DataBlock, encodings::{ - logical::primitive::miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + logical::primitive::miniblock::{ + MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, + }, physical::block::{CompressionConfig, GeneralBufferCompressor}, }, format::{ProtobufUtils21, pb21::CompressiveEncoding}, }; +use lance_core::Error; + +pub(crate) fn compress_block( + compression: CompressionConfig, + payload: &[u8], +) -> Result { + let compressor = GeneralBufferCompressor::get_compressor(compression)?; + let mut compressed = Vec::new(); + compressor.compress(payload, &mut compressed)?; + Ok(LanceBuffer::from(compressed)) +} + +pub(crate) fn decompress_block_exact( + compression: CompressionConfig, + payload: &LanceBuffer, + expected_bytes: usize, +) -> Result { + let compressor = GeneralBufferCompressor::get_compressor(compression)?; + let mut decompressed = Vec::new(); + compressor.decompress_exact(payload, &mut decompressed, expected_bytes)?; + Ok(LanceBuffer::from(decompressed)) +} + +/// General-purpose block compressor that owns its child block compressor. +#[derive(Debug)] +pub(crate) struct GeneralBlockCompressor { + child: Box, + compression: CompressionConfig, +} + +impl GeneralBlockCompressor { + pub(crate) fn new(child: Box, compression: CompressionConfig) -> Self { + Self { child, compression } + } +} + +impl BlockCompressor for GeneralBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let payload = self.child.compress(data)?.ok_or_else(|| { + Error::invalid_input("General block compression requires a payload-bearing child") + })?; + compress_block(self.compression, &payload).map(Some) + } +} + +/// Bounded general block decompressor used by generic fixed-width sequences. +#[derive(Debug)] +pub(crate) struct GenericGeneralBlockDecompressor { + child: Box, + compression: CompressionConfig, + value_type: BlockValueType, +} + +impl GenericGeneralBlockDecompressor { + pub(crate) fn new( + child: Box, + compression: CompressionConfig, + value_type: BlockValueType, + ) -> Self { + Self { + child, + compression, + value_type, + } + } +} + +impl BlockDecompressor for GenericGeneralBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "General block compression")?; + let expected_bytes = usize::try_from(num_values) + .ok() + .and_then(|values| values.checked_mul(self.value_type.bits_per_value() as usize / 8)) + .ok_or_else(|| Error::invalid_input("General block output length overflows usize"))?; + let decompressed = decompress_block_exact(self.compression, &data, expected_bytes)?; + self.child.decompress(Some(decompressed), num_values) + } +} /// A miniblock compressor that wraps another miniblock compressor and applies /// general-purpose compression (LZ4, Zstd) to the resulting buffers. @@ -27,18 +110,12 @@ impl GeneralMiniBlockCompressor { pub fn new(inner: Box, compression: CompressionConfig) -> Self { Self { inner, compression } } -} - -/// Minimum buffer size to consider for compression -const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024; - -use super::super::logical::primitive::miniblock::MiniBlockChunk; - -impl MiniBlockCompressor for GeneralMiniBlockCompressor { - fn compress(&self, page: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { - // First, compress with the inner compressor - let (inner_compressed, inner_encoding) = self.inner.compress(page)?; + fn compress_inner( + &self, + inner_compressed: MiniBlockCompressed, + inner_encoding: CompressiveEncoding, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { // Return the original encoding without compression if there's no data or // the first buffer is not large enough if inner_compressed.data.is_empty() @@ -108,6 +185,22 @@ impl MiniBlockCompressor for GeneralMiniBlockCompressor { } } +/// Minimum buffer size to consider for compression +const MIN_BUFFER_SIZE_FOR_COMPRESSION: usize = 4 * 1024; + +use super::super::logical::primitive::miniblock::MiniBlockChunk; + +impl MiniBlockCompressor for GeneralMiniBlockCompressor { + fn compress( + &self, + page: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + let (inner_compressed, inner_encoding) = self.inner.compress(page, context)?; + self.compress_inner(inner_compressed, inner_encoding) + } +} + /// A miniblock decompressor that first decompresses buffers using general-purpose /// compression (LZ4, Zstd) and then delegates to an inner miniblock decompressor. #[derive(Debug)] @@ -146,6 +239,10 @@ mod tests { use crate::format::pb21::compressive_encoding::Compression; use arrow_array::{Float64Array, Int32Array}; + fn miniblock_context() -> MiniBlockCompressionContext { + MiniBlockCompressionContext::new(0, true, true) + } + #[derive(Debug)] struct TestCase { name: &'static str, @@ -249,7 +346,9 @@ mod tests { GeneralMiniBlockCompressor::new(test_case.inner_encoder, test_case.compression); // Compress the data - let (compressed, encoding) = compressor.compress(test_case.data).unwrap(); + let (compressed, encoding) = compressor + .compress(test_case.data, miniblock_context()) + .unwrap(); // Check if compression was applied as expected match &encoding.compression { @@ -461,7 +560,7 @@ mod tests { let compressor = GeneralMiniBlockCompressor::new(inner, compression); // Compress the data - let (compressed, encoding) = compressor.compress(block).unwrap(); + let (compressed, encoding) = compressor.compress(block, miniblock_context()).unwrap(); // Should get GeneralMiniBlock encoding since buffer is 4KB match &encoding.compression { @@ -503,7 +602,7 @@ mod tests { }, ); - let (compressed, _) = compressor.compress(data).unwrap(); + let (compressed, _) = compressor.compress(data, miniblock_context()).unwrap(); // RLE produces 2 buffers, but only the first one is compressed assert_eq!(compressed.data.len(), 2); } @@ -539,7 +638,9 @@ mod tests { }, ); - let (_compressed, encoding) = compressor.compress(test_32.data).unwrap(); + let (_compressed, encoding) = compressor + .compress(test_32.data, miniblock_context()) + .unwrap(); // Verify the encoding structure match &encoding.compression { @@ -596,7 +697,9 @@ mod tests { }, ); - let (_compressed_64, encoding_64) = compressor_64.compress(block_64).unwrap(); + let (_compressed_64, encoding_64) = compressor_64 + .compress(block_64, miniblock_context()) + .unwrap(); // Verify the encoding structure for 64-bit match &encoding_64.compression { @@ -650,7 +753,7 @@ mod tests { }, ); - let result = compressor.compress(empty_block); + let result = compressor.compress(empty_block, miniblock_context()); match result { Ok((compressed, _)) => { assert_eq!(compressed.num_values, 0); diff --git a/rust/lance-encoding/src/encodings/physical/packed.rs b/rust/lance-encoding/src/encodings/physical/packed.rs index ad2221dffed..96aa61e69d6 100644 --- a/rust/lance-encoding/src/encodings/physical/packed.rs +++ b/rust/lance-encoding/src/encodings/physical/packed.rs @@ -27,7 +27,7 @@ use crate::{ }, encodings::logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::{MiniBlockCompressed, MiniBlockCompressor}, + miniblock::{MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor}, }, format::{ ProtobufUtils21, @@ -73,7 +73,11 @@ fn struct_data_block_to_fixed_width_data_block( pub struct PackedStructFixedWidthMiniBlockEncoder {} impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::Struct(struct_data_block) => { let bits_per_values = struct_data_block.children.iter().map(|data_block| data_block.as_fixed_width_ref().unwrap().bits_per_value).collect::>(); @@ -84,7 +88,7 @@ impl MiniBlockCompressor for PackedStructFixedWidthMiniBlockEncoder { // store and transformed fixed-width data block. let value_miniblock_compressor = Box::new(ValueEncoder::default()) as Box; let (value_miniblock_compressed, value_array_encoding) = - value_miniblock_compressor.compress(data_block)?; + value_miniblock_compressor.compress(data_block, context)?; Ok(( value_miniblock_compressed, diff --git a/rust/lance-encoding/src/encodings/physical/range.rs b/rust/lance-encoding/src/encodings/physical/range.rs new file mode 100644 index 00000000000..21a06c111ec --- /dev/null +++ b/rust/lance-encoding/src/encodings/physical/range.rs @@ -0,0 +1,245 @@ +// 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; +use crate::{ + buffer::LanceBuffer, + compression::{BlockCompressor, 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. +#[derive(Debug)] +pub(crate) struct RangeEncoder { + bits_per_value: u64, + start: u64, + step: u64, +} + +impl RangeEncoder { + pub(crate) fn new(bits_per_value: u64, start: u64, step: u64) -> Self { + Self { + bits_per_value, + start, + step, + } + } +} + +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(), + })) +} + +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 b04716c2b44..0e7a9ecf58e 100644 --- a/rust/lance-encoding/src/encodings/physical/rle.rs +++ b/rust/lance-encoding/src/encodings/physical/rle.rs @@ -58,14 +58,24 @@ use arrow_buffer::{ArrowNativeType, ScalarBuffer}; use log::trace; use crate::buffer::LanceBuffer; -use crate::compression::{BlockCompressor, BlockDecompressor, MiniBlockDecompressor}; +use crate::compression::{ + BlockCompressor, BlockDecompressor, BlockValueType, MiniBlockDecompressor, + block::{ + fixed_block, fixed_from_u64_values, read_unsigned_values, validate_fixed_payload_len, + visit_unsigned_values, + }, + require_block_payload, +}; use crate::data::DataBlock; use crate::data::{BlockInfo, FixedWidthDataBlock}; use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, - MiniBlockCompressor, + MiniBlockCompressionContext, MiniBlockCompressor, +}; +use crate::encodings::physical::{ + block::{BufferCompressor, CompressionConfig, GeneralBufferCompressor}, + try_vec_with_capacity, }; -use crate::encodings::physical::block::{CompressionConfig, GeneralBufferCompressor}; use crate::format::ProtobufUtils21; use crate::format::pb21::CompressiveEncoding; @@ -872,7 +882,9 @@ impl RleEncoder { num_values: child_values, block_info: BlockInfo::default(), }); - let chunk_packed = BlockCompressor::compress(&compressor, block)?; + let chunk_packed = BlockCompressor::compress(&compressor, block)?.ok_or_else(|| { + Error::internal("RLE bitpacking child returned no payload".to_string()) + })?; let packed_size = u32::try_from(chunk_packed.len()).map_err(|_| { Error::invalid_input_source( format!( @@ -1036,7 +1048,11 @@ impl RleChildCandidate { } impl MiniBlockCompressor for RleEncoder { - fn compress(&self, data: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + data: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match data { DataBlock::FixedWidth(fixed_width) => { let num_values = fixed_width.num_values; @@ -1103,7 +1119,7 @@ impl MiniBlockCompressor for RleEncoder { impl BlockCompressor for RleEncoder { // Block format: [8-byte header: values buffer size][values buffer][run_lengths buffer] - fn compress(&self, data: DataBlock) -> Result { + fn compress(&self, data: DataBlock) -> Result> { match data { DataBlock::FixedWidth(fixed_width) => { let num_values = fixed_width.num_values; @@ -1118,7 +1134,7 @@ impl BlockCompressor for RleEncoder { combined.extend_from_slice(&values_size.to_le_bytes()); combined.extend_from_slice(&all_buffers[0]); combined.extend_from_slice(&all_buffers[1]); - Ok(LanceBuffer::from(combined)) + Ok(Some(LanceBuffer::from(combined))) } _ => Err(Error::invalid_input_source( "RLE encoding only supports FixedWidth data blocks".into(), @@ -1145,6 +1161,9 @@ pub(crate) struct RleChildDecompressor { #[derive(Debug)] enum RleChildDecompressorInner { Flat, + General { + compressor: Box, + }, Block { decompressor: Box, requires_num_values: bool, @@ -1173,13 +1192,22 @@ impl RleChildDecompressor { } } + pub(crate) fn general(bits_per_value: u64, compression: CompressionConfig) -> Result { + Ok(Self { + bits_per_value, + inner: RleChildDecompressorInner::General { + compressor: GeneralBufferCompressor::get_compressor(compression)?, + }, + }) + } + pub(crate) fn bits_per_value(&self) -> u64 { self.bits_per_value } pub(crate) fn requires_num_values(&self) -> bool { match &self.inner { - RleChildDecompressorInner::Flat => false, + RleChildDecompressorInner::Flat | RleChildDecompressorInner::General { .. } => false, RleChildDecompressorInner::Block { requires_num_values, .. @@ -1195,10 +1223,66 @@ impl RleChildDecompressor { &self, data: LanceBuffer, num_values: Option, + max_num_values: u64, label: &str, ) -> Result { match &self.inner { RleChildDecompressorInner::Flat => Ok(data), + RleChildDecompressorInner::General { compressor } => { + let bytes_per_value = usize::try_from(self.bits_per_value / 8).map_err(|_| { + Error::invalid_input_source( + format!( + "RLE {label} child bit width is too large: {}", + self.bits_per_value + ) + .into(), + ) + })?; + let max_output_bytes = usize::try_from(max_num_values) + .ok() + .and_then(|num_values| num_values.checked_mul(bytes_per_value)) + .ok_or_else(|| { + Error::invalid_input_source( + format!("RLE {label} child maximum payload length overflows usize") + .into(), + ) + })?; + let mut decompressed = Vec::new(); + if let Some(num_values) = num_values { + if num_values > max_num_values { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child expects {num_values} values, exceeding the maximum run count {max_num_values}" + ) + .into(), + )); + } + let expected_bytes = usize::try_from(num_values) + .ok() + .and_then(|num_values| num_values.checked_mul(bytes_per_value)) + .ok_or_else(|| { + Error::invalid_input_source( + format!( + "RLE {label} child expected payload length overflows usize" + ) + .into(), + ) + })?; + compressor.decompress_exact(&data, &mut decompressed, expected_bytes)?; + } else { + compressor.decompress_bounded(&data, &mut decompressed, max_output_bytes)?; + } + if bytes_per_value == 0 || !decompressed.len().is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input_source( + format!( + "RLE {label} child decompressed to {} bytes, not divisible by {bytes_per_value}", + decompressed.len() + ) + .into(), + )); + } + Ok(LanceBuffer::from(decompressed)) + } RleChildDecompressorInner::Block { decompressor, requires_num_values, @@ -1212,7 +1296,7 @@ impl RleChildDecompressor { } else { num_values.unwrap_or(0) }; - let decoded = decompressor.decompress(data, num_values)?; + let decoded = decompressor.decompress(Some(data), num_values)?; self.extract_fixed_width(decoded, num_values, label) } } @@ -1318,7 +1402,7 @@ impl RleDecompressor { let values_buffer = data_iter.next().unwrap(); let lengths_buffer = data_iter.next().unwrap(); let (values_buffer, lengths_buffer) = - self.decode_child_buffers(values_buffer, lengths_buffer)?; + self.decode_child_buffers(values_buffer, lengths_buffer, num_values)?; let decoded_data = match self.bits_per_value { 8 => self.decode_generic::( @@ -1368,6 +1452,7 @@ impl RleDecompressor { &self, values_buffer: LanceBuffer, lengths_buffer: LanceBuffer, + max_num_runs: u64, ) -> Result<(LanceBuffer, LanceBuffer)> { let values_requires_num_runs = self.values.requires_num_values(); let lengths_requires_num_runs = self.run_lengths.requires_num_values(); @@ -1378,31 +1463,38 @@ impl RleDecompressor { } if values_requires_num_runs { - let lengths_buffer = self - .run_lengths - .decode(lengths_buffer, None, "run lengths")?; + let lengths_buffer = + self.run_lengths + .decode(lengths_buffer, None, max_num_runs, "run lengths")?; let num_runs = Self::num_child_values( &lengths_buffer, self.run_lengths.bits_per_value(), "run lengths", )?; - let values_buffer = self - .values - .decode(values_buffer, Some(num_runs), "values")?; + let values_buffer = + self.values + .decode(values_buffer, Some(num_runs), max_num_runs, "values")?; Ok((values_buffer, lengths_buffer)) } else if lengths_requires_num_runs { - let values_buffer = self.values.decode(values_buffer, None, "values")?; + let values_buffer = self + .values + .decode(values_buffer, None, max_num_runs, "values")?; let num_runs = Self::num_child_values(&values_buffer, self.values.bits_per_value(), "values")?; - let lengths_buffer = - self.run_lengths - .decode(lengths_buffer, Some(num_runs), "run lengths")?; + let lengths_buffer = self.run_lengths.decode( + lengths_buffer, + Some(num_runs), + max_num_runs, + "run lengths", + )?; Ok((values_buffer, lengths_buffer)) } else { - let values_buffer = self.values.decode(values_buffer, None, "values")?; - let lengths_buffer = self - .run_lengths - .decode(lengths_buffer, None, "run lengths")?; + let values_buffer = self + .values + .decode(values_buffer, None, max_num_runs, "values")?; + let lengths_buffer = + self.run_lengths + .decode(lengths_buffer, None, max_num_runs, "run lengths")?; Ok((values_buffer, lengths_buffer)) } } @@ -1555,7 +1647,8 @@ impl MiniBlockDecompressor for RleDecompressor { } impl BlockDecompressor for RleDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "RLE")?; let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?; self.decode_data(vec![values_buffer, lengths_buffer], num_values, false) } @@ -1825,7 +1918,7 @@ impl RleDecompressor { } let (values_buffer, lengths_buffer) = parse_rle_block_frame(&data)?; let (values_buffer, lengths_buffer) = - self.decode_child_buffers(values_buffer, lengths_buffer)?; + self.decode_child_buffers(values_buffer, lengths_buffer, num_values)?; RleRuns::try_new( values_buffer, lengths_buffer, @@ -1835,16 +1928,451 @@ impl RleDecompressor { } } +// Generic block codec support. + +pub(crate) const BLOCK_FRAME_BYTES: u64 = 8; + +/// Generic RLE compressor that owns both child compressors. +#[derive(Debug)] +pub(crate) struct BlockRleCompressor { + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, +} + +impl BlockRleCompressor { + pub(crate) fn new( + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, + ) -> Self { + Self { + value_type, + run_length_type, + values, + run_lengths, + } + } +} + +impl BlockCompressor for BlockRleCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "RLE block compression requires fixed-width data", + )); + }; + let (run_values, run_lengths) = + materialize_block(&data, self.value_type, self.run_length_type)?; + let values_payload = self.values.compress(DataBlock::FixedWidth(run_values))?; + let lengths_payload = self + .run_lengths + .compress(DataBlock::FixedWidth(run_lengths))?; + + if values_payload.is_none() && lengths_payload.is_none() { + return Ok(None); + } + + let values_payload = values_payload.unwrap_or_else(LanceBuffer::empty); + let lengths_payload = lengths_payload.unwrap_or_else(LanceBuffer::empty); + let mut output = try_block_frame(values_payload.len(), lengths_payload.len())?; + output.extend_from_slice(&(values_payload.len() as u64).to_le_bytes()); + output.extend_from_slice(&values_payload); + output.extend_from_slice(&lengths_payload); + Ok(Some(LanceBuffer::from(output))) + } +} + +fn try_block_frame(values_payload_bytes: usize, lengths_payload_bytes: usize) -> Result> { + let capacity = (BLOCK_FRAME_BYTES as usize) + .checked_add(values_payload_bytes) + .and_then(|capacity| capacity.checked_add(lengths_payload_bytes)) + .ok_or_else(|| Error::invalid_input("RLE frame length overflows usize"))?; + let mut output = Vec::new(); + output.try_reserve_exact(capacity).map_err(|error| { + Error::invalid_input(format!( + "RLE could not reserve {capacity} frame bytes: {error}" + )) + })?; + Ok(output) +} + +/// Metadata-only run lengths used to recover the RLE child cardinality. +#[derive(Debug, Clone, Copy)] +pub(crate) enum MetadataRunLengths { + Constant(u64), + Range { start: u64, step: u64 }, +} + +impl MetadataRunLengths { + fn infer_run_count(self, num_values: u64) -> Result { + match self { + Self::Constant(run_length) => { + if run_length == 0 { + return Err(Error::invalid_input("RLE run lengths must be positive")); + } + if !num_values.is_multiple_of(run_length) { + return Err(Error::invalid_input(format!( + "RLE constant run length {run_length} does not divide {num_values} values" + ))); + } + let run_count = num_values / run_length; + if run_count == 0 { + return Err(Error::invalid_input( + "RLE metadata describes zero runs for a non-empty sequence", + )); + } + Ok(run_count) + } + 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}" + ))) + } + } + } +} + +/// The descriptor-derived source of the RLE run count. +#[derive(Debug, Clone, Copy)] +pub(crate) enum BlockRunCount { + Metadata(MetadataRunLengths), + ValuesPayload, + RunLengthsPayload, +} + +/// Generic RLE decompressor that owns both child decompressors. +#[derive(Debug)] +pub(crate) struct BlockRleDecompressor { + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, + values_have_payload: bool, + run_lengths_have_payload: bool, + run_count: BlockRunCount, +} + +impl BlockRleDecompressor { + pub(crate) fn new( + value_type: BlockValueType, + run_length_type: BlockValueType, + values: Box, + run_lengths: Box, + values_have_payload: bool, + run_lengths_have_payload: bool, + run_count: BlockRunCount, + ) -> Self { + Self { + value_type, + run_length_type, + values, + run_lengths, + values_have_payload, + run_lengths_have_payload, + run_count, + } + } +} + +impl BlockDecompressor for BlockRleDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + if num_values == 0 { + return Err(Error::invalid_input( + "RLE cannot represent an empty block sequence", + )); + } + let (values_payload, lengths_payload) = if self.values_have_payload + || self.run_lengths_have_payload + { + let data = require_block_payload(data, "RLE")?; + if data.len() < BLOCK_FRAME_BYTES as usize { + return Err(Error::invalid_input(format!( + "RLE payload has {} bytes, shorter than its {BLOCK_FRAME_BYTES}-byte header", + data.len() + ))); + } + let values_size = u64::from_le_bytes( + data[..BLOCK_FRAME_BYTES as usize] + .try_into() + .expect("RLE header length was checked"), + ); + let values_size = usize::try_from(values_size).map_err(|_| { + Error::invalid_input("RLE values payload length does not fit usize") + })?; + let values_start = BLOCK_FRAME_BYTES as usize; + let lengths_start = values_start + .checked_add(values_size) + .ok_or_else(|| Error::invalid_input("RLE values payload end overflows usize"))?; + if lengths_start > data.len() { + return Err(Error::invalid_input(format!( + "RLE values payload ends at {lengths_start}, beyond {} bytes", + data.len() + ))); + } + let values_payload = data.slice_with_length(values_start, values_size); + let lengths_payload = data.slice_with_length(lengths_start, data.len() - lengths_start); + if !self.values_have_payload && !values_payload.is_empty() { + return Err(Error::invalid_input(format!( + "Metadata-only RLE values child has {} framed payload bytes", + values_payload.len() + ))); + } + if !self.run_lengths_have_payload && !lengths_payload.is_empty() { + return Err(Error::invalid_input(format!( + "Metadata-only RLE run-length child has {} framed payload bytes", + lengths_payload.len() + ))); + } + ( + self.values_have_payload.then_some(values_payload), + self.run_lengths_have_payload.then_some(lengths_payload), + ) + } else { + if data.is_some() { + return Err(Error::invalid_input("Metadata-only RLE expects no payload")); + } + (None, None) + }; + + let run_count = match self.run_count { + BlockRunCount::Metadata(metadata) => metadata.infer_run_count(num_values)?, + BlockRunCount::ValuesPayload => infer_flat_run_count( + values_payload.as_ref().ok_or_else(|| { + Error::invalid_input("RLE values payload is required to infer the run count") + })?, + self.value_type, + num_values, + "RLE values", + )?, + BlockRunCount::RunLengthsPayload => infer_flat_run_count( + lengths_payload.as_ref().ok_or_else(|| { + Error::invalid_input( + "RLE run-length payload is required to infer the run count", + ) + })?, + self.run_length_type, + num_values, + "RLE run lengths", + )?, + }; + if run_count == 0 { + return Err(Error::invalid_input( + "RLE payload contains zero runs for a non-empty sequence", + )); + } + let values = self.values.decompress(values_payload, run_count)?; + let run_lengths = self.run_lengths.decompress(lengths_payload, run_count)?; + expand_block( + values, + run_lengths, + self.value_type, + self.run_length_type, + num_values, + run_count, + ) + } +} + +fn infer_flat_run_count( + payload: &LanceBuffer, + value_type: BlockValueType, + max_run_count: u64, + label: &str, +) -> Result { + let bytes_per_value = value_type.bytes_per_value(); + if !payload.len().is_multiple_of(bytes_per_value) { + return Err(Error::invalid_input(format!( + "{label} payload has {} bytes, not divisible by {bytes_per_value}", + payload.len() + ))); + } + let run_count = (payload.len() / bytes_per_value) as u64; + if run_count > max_run_count { + return Err(Error::invalid_input(format!( + "{label} payload contains {run_count} runs, exceeding the {max_run_count}-run limit" + ))); + } + Ok(run_count) +} + +pub(crate) fn expand_block( + values: DataBlock, + run_lengths: DataBlock, + value_type: BlockValueType, + run_length_type: BlockValueType, + num_values: u64, + run_count: u64, +) -> Result { + let DataBlock::FixedWidth(values) = values else { + return Err(Error::invalid_input( + "RLE values decoded to a non fixed-width block", + )); + }; + let DataBlock::FixedWidth(run_lengths) = run_lengths else { + return Err(Error::invalid_input( + "RLE run lengths decoded to a non fixed-width block", + )); + }; + if values.num_values != run_count + || values.bits_per_value != value_type.bits_per_value() + || run_lengths.num_values != run_count + || run_lengths.bits_per_value != run_length_type.bits_per_value() + { + return Err(Error::invalid_input( + "RLE child cardinality or bit width does not match its descriptor", + )); + } + validate_fixed_payload_len(&values.data, value_type, run_count, "RLE values")?; + validate_fixed_payload_len( + &run_lengths.data, + run_length_type, + run_count, + "RLE run lengths", + )?; + + let lengths = read_unsigned_values(&run_lengths, run_length_type)?; + let mut total = 0_u64; + for (index, length) in lengths.iter().enumerate() { + if *length == 0 { + return Err(Error::invalid_input(format!( + "RLE run length at index {index} is zero" + ))); + } + total = total.checked_add(*length).ok_or_else(|| { + Error::invalid_input(format!("RLE run length sum overflows at index {index}")) + })?; + } + if total != num_values { + return Err(Error::invalid_input(format!( + "RLE run lengths sum to {total}, expected {num_values}" + ))); + } + let lengths = lengths + .into_iter() + .map(|length| { + usize::try_from(length) + .map_err(|_| Error::invalid_input("RLE run length does not fit usize")) + }) + .collect::>>()?; + + let run_values = read_unsigned_values(&values, value_type)?; + let output = match value_type { + BlockValueType::UInt8 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value as u8, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt16 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value as u16, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt32 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value as u32, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + BlockValueType::UInt64 => { + let mut output = try_vec_with_capacity::(num_values, "RLE output")?; + for (value, length) in run_values.iter().zip(&lengths) { + output.extend(std::iter::repeat_n(*value, *length)); + } + LanceBuffer::reinterpret_vec(output) + } + }; + Ok(fixed_block(value_type, num_values, output)) +} + +pub(crate) fn materialize_block( + data: &FixedWidthDataBlock, + value_type: BlockValueType, + run_length_type: BlockValueType, +) -> Result<(FixedWidthDataBlock, FixedWidthDataBlock)> { + if data.num_values == 0 { + return Err(Error::invalid_input( + "RLE cannot materialize an empty sequence", + )); + } + let max_run_length = run_length_type.max_value(); + let mut run_values = Vec::new(); + let mut run_lengths = Vec::new(); + let mut current = None; + let mut length = 0_u64; + visit_unsigned_values(data, value_type, |value| { + match current { + Some(current_value) if value == current_value && length < max_run_length => { + length += 1; + } + Some(current_value) => { + run_values.push(current_value); + run_lengths.push(length); + current = Some(value); + length = 1; + } + None => { + current = Some(value); + length = 1; + } + } + Ok(()) + })?; + run_values.push(current.expect("non-empty input has a run value")); + run_lengths.push(length); + Ok(( + fixed_from_u64_values(&run_values, value_type, "RLE values")?, + fixed_from_u64_values(&run_lengths, run_length_type, "RLE run lengths")?, + )) +} + #[cfg(test)] mod tests { use std::sync::Arc; use super::*; - use crate::compression::{ - DecompressionStrategy, DefaultDecompressionStrategy, create_rle_decompressor, - }; + #[cfg(any(feature = "lz4", feature = "zstd"))] + use crate::compression::create_rle_decompressor; + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] + use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; use crate::data::DataBlock; use crate::encodings::logical::primitive::miniblock::MAX_MINIBLOCK_VALUES; + #[cfg(any(feature = "lz4", feature = "zstd"))] use crate::encodings::physical::block::{CompressionConfig, CompressionScheme}; use crate::{ buffer::LanceBuffer, @@ -1875,11 +2403,16 @@ mod tests { num_values, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); - let eager = - BlockDecompressor::decompress(&RleDecompressor::new(16), frame.clone(), num_values) - .unwrap(); + let eager = BlockDecompressor::decompress( + &RleDecompressor::new(16), + Some(frame.clone()), + num_values, + ) + .unwrap(); let DataBlock::FixedWidth(eager) = eager else { panic!("expected fixed-width block"); }; @@ -1952,7 +2485,9 @@ mod tests { num_values, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); let (values, lengths) = parse_rle_block_frame(&frame).unwrap(); let compression = test_general_compression(); @@ -1994,7 +2529,9 @@ mod tests { num_values, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); let runs = RleDecompressor::new(16) .decode_u16_runs(frame, num_values) .unwrap(); @@ -2015,7 +2552,9 @@ mod tests { num_values: n, block_info: BlockInfo::new(), }); - let frame = BlockCompressor::compress(&RleEncoder::new(), block).unwrap(); + let frame = BlockCompressor::compress(&RleEncoder::new(), block) + .unwrap() + .unwrap(); let runs = RleDecompressor::new(16).decode_u16_runs(frame, n).unwrap(); assert_eq!( runs.coalesced_runs() as u64, @@ -2025,6 +2564,19 @@ mod tests { assert_eq!(expand_u16_runs(&runs), alternating); } + fn compress_miniblock( + compressor: &dyn MiniBlockCompressor, + data: DataBlock, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) + } + + #[test] + fn block_frame_capacity_overflow_is_fallible() { + let error = try_block_frame(usize::MAX, 1).unwrap_err(); + assert!(error.to_string().contains("frame length overflows usize")); + } + // ========== Core Functionality Tests ========== #[test] @@ -2035,7 +2587,7 @@ mod tests { let array = Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 3, 3]); let data_block = DataBlock::from_array(array); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, data_block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, data_block).unwrap(); assert_eq!(compressed.num_values, 9); assert_eq!(compressed.chunks.len(), 1); @@ -2056,8 +2608,7 @@ mod tests { data.extend(&[100i32; 300]); // Will be split into 255+45 let array = Int32Array::from(data); - let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Should have 6 runs total (4 for first value, 2 for second) let lengths_buffer = &compressed.data[1]; @@ -2071,7 +2622,7 @@ mod tests { let data = vec![42i32; 1000]; let array = Int32Array::from(data); let (compressed, encoding) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); assert_eq!(compressed.data[0].len(), 4); assert_eq!(compressed.data[1].len(), 2); @@ -2112,7 +2663,7 @@ mod tests { RleEncoder::with_child_encoding(RunLengthWidth::U8, Some(compression), None, false); let array = Int32Array::from(repeating_runs(1024, 4)); let (compressed, encoding) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); let rle = expect_rle(&encoding); assert!(matches!( @@ -2145,7 +2696,7 @@ mod tests { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, Some(compression), false); let expected = repeating_runs(1024, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2181,7 +2732,7 @@ mod tests { use crate::encodings::physical::bitpacking::OutOfLineBitpacking; let expected = repeating_runs(1024, 4); - let (compressed, _) = MiniBlockCompressor::compress( + let (compressed, _) = compress_miniblock( &RleEncoder::new(), DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2195,7 +2746,9 @@ mod tests { block_info: BlockInfo::default(), }); let bitpacked_run_lengths = - BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block).unwrap(); + BlockCompressor::compress(&OutOfLineBitpacking::new(3, 8), run_lengths_block) + .unwrap() + .unwrap(); let encoding = ProtobufUtils21::rle( ProtobufUtils21::flat(32, None), ProtobufUtils21::out_of_line_bitpacking(8, ProtobufUtils21::flat(3, None)), @@ -2239,6 +2792,7 @@ mod tests { } } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn repeating_runs(num_runs: usize, run_length: usize) -> Vec { let mut values = Vec::with_capacity(num_runs * run_length); for run in 0..num_runs { @@ -2247,6 +2801,7 @@ mod tests { values } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn expect_rle(encoding: &CompressiveEncoding) -> &crate::format::pb21::Rle { match encoding.compression.as_ref().unwrap() { crate::format::pb21::compressive_encoding::Compression::Rle(rle) => rle, @@ -2254,6 +2809,7 @@ mod tests { } } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn assert_decoded_i32_eq(decoded: DataBlock, expected: &[i32]) { match decoded { DataBlock::FixedWidth(block) => { @@ -2275,7 +2831,7 @@ mod tests { false, ); let expected = repeating_runs(8192, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2306,7 +2862,7 @@ mod tests { fn test_rle_miniblock_bitpacks_values_child_when_smaller() { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); let expected = monotonic_runs(2048, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2336,7 +2892,7 @@ mod tests { fn test_rle_miniblock_bitpacks_run_lengths_when_values_do_not_shrink() { let encoder = RleEncoder::with_child_encoding(RunLengthWidth::U8, None, None, true); let expected = high_entropy_runs(2048, 4); - let (compressed, encoding) = MiniBlockCompressor::compress( + let (compressed, encoding) = compress_miniblock( &encoder, DataBlock::from_array(Int32Array::from(expected.clone())), ) @@ -2361,6 +2917,7 @@ mod tests { assert_eq!(decoded, expected); } + #[cfg(any(feature = "bitpacking", feature = "lz4", feature = "zstd"))] fn decompress_i32_chunks( compressed: &MiniBlockCompressed, encoding: &CompressiveEncoding, @@ -2464,7 +3021,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); let decompressor = RleDecompressor::new(bits_per_value); let decompressed = MiniBlockDecompressor::decompress( &decompressor, @@ -2498,7 +3055,7 @@ mod tests { let array = Int32Array::from(data); let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Verify all non-last chunks have power-of-2 values for (i, chunk) in compressed.chunks.iter().enumerate() { @@ -2683,8 +3240,9 @@ mod tests { payload.extend_from_slice(&values); payload.extend_from_slice(&lengths); - let error = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(payload), 5) - .unwrap_err(); + let error = + BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(payload)), 5) + .unwrap_err(); assert!(matches!(&error, Error::InvalidInput { .. })); assert!( error @@ -2737,7 +3295,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, empty_block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, empty_block).unwrap(); assert_eq!(compressed.num_values, 0); assert!(compressed.data.is_empty()); @@ -2771,8 +3329,7 @@ mod tests { data.extend(vec![777i32; 2000]); let array = Int32Array::from(data.clone()); - let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Manually decompress all chunks let mut reconstructed = Vec::new(); @@ -2885,7 +3442,7 @@ mod tests { // Compress the data let array = Int32Array::from(data.clone()); let (compressed, _) = - MiniBlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + compress_miniblock(&encoder, DataBlock::from_array(array)).unwrap(); // Decompress and verify match MiniBlockDecompressor::decompress( @@ -2955,7 +3512,7 @@ mod tests { block_info: BlockInfo::default(), }); - let (compressed, _) = MiniBlockCompressor::compress(&encoder, block).unwrap(); + let (compressed, _) = compress_miniblock(&encoder, block).unwrap(); // Debug first few chunks for (i, chunk) in compressed.chunks.iter().take(5).enumerate() { @@ -3131,13 +3688,97 @@ mod tests { } // ========== Block Related tests ========== + #[cfg(any(feature = "lz4", feature = "zstd"))] + fn assert_rle_general_children_reject_oversized_output( + compression: CompressionConfig, + payload: Vec, + ) { + let max_num_values = 2; + + let values_general = RleChildDecompressor::general(32, compression).unwrap(); + let values_decompressor = RleDecompressor::with_child_decompressors( + 32, + RunLengthWidth::U8, + values_general, + RleChildDecompressor::flat(8), + ); + let error = MiniBlockDecompressor::decompress( + &values_decompressor, + vec![ + LanceBuffer::from(payload.clone()), + LanceBuffer::from(vec![2_u8]), + ], + max_num_values, + ) + .unwrap_err(); + assert!( + error.to_string().contains("exceed"), + "unexpected values-child error: {error}" + ); + + let lengths_general = RleChildDecompressor::general(8, compression).unwrap(); + let lengths_decompressor = RleDecompressor::with_child_decompressors( + 32, + RunLengthWidth::U8, + RleChildDecompressor::flat(32), + lengths_general, + ); + let error = MiniBlockDecompressor::decompress( + &lengths_decompressor, + vec![ + LanceBuffer::reinterpret_vec(vec![7_u32]), + LanceBuffer::from(payload), + ], + max_num_values, + ) + .unwrap_err(); + assert!( + error.to_string().contains("exceed"), + "unexpected run-lengths-child error: {error}" + ); + } + + #[cfg(feature = "lz4")] + #[test] + fn test_rle_general_children_bound_lz4_output() { + assert_rle_general_children_reject_oversized_output( + CompressionConfig::new(CompressionScheme::Lz4, None), + u32::MAX.to_le_bytes().to_vec(), + ); + } + + #[cfg(feature = "zstd")] + #[test] + fn test_rle_general_children_bound_length_prefixed_zstd_output() { + assert_rle_general_children_reject_oversized_output( + CompressionConfig::new(CompressionScheme::Zstd, Some(0)), + u64::MAX.to_le_bytes().to_vec(), + ); + } + + #[cfg(feature = "zstd")] + #[test] + fn test_rle_general_children_bound_raw_zstd_output() { + use std::io::Write; + + let mut payload = Vec::new(); + let mut encoder = ::zstd::Encoder::new(&mut payload, 0).unwrap(); + encoder.write_all(&[0; 1024]).unwrap(); + encoder.finish().unwrap(); + + assert_rle_general_children_reject_oversized_output( + CompressionConfig::new(CompressionScheme::Zstd, Some(0)), + payload, + ); + } + #[test] fn test_block_decompressor_rejects_overflowing_values_size() { let decompressor = RleDecompressor::new(32); let mut data = Vec::new(); data.extend_from_slice(&u64::MAX.to_le_bytes()); - let result = BlockDecompressor::decompress(&decompressor, LanceBuffer::from(data), 1); + let result = BlockDecompressor::decompress(&decompressor, Some(LanceBuffer::from(data)), 1); assert!(result.is_err()); assert!( result @@ -3150,8 +3791,11 @@ mod tests { #[test] fn test_block_decompressor_too_small() { let decompressor = RleDecompressor::new(32); - let result = - BlockDecompressor::decompress(&decompressor, LanceBuffer::from(vec![1, 2, 3]), 10); + let result = BlockDecompressor::decompress( + &decompressor, + Some(LanceBuffer::from(vec![1, 2, 3])), + 10, + ); assert!(result.is_err()); assert!( result @@ -3167,7 +3811,9 @@ mod tests { let data = vec![1i32, 1, 1]; let array = Int32Array::from(data); - let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)).unwrap(); + let compressed = BlockCompressor::compress(&encoder, DataBlock::from_array(array)) + .unwrap() + .unwrap(); // Verify header format: first 8 bytes should be values_size as u64 assert!(compressed.len() >= 8); diff --git a/rust/lance-encoding/src/encodings/physical/value.rs b/rust/lance-encoding/src/encodings/physical/value.rs index 606f49b699a..7789e7cfe3d 100644 --- a/rust/lance-encoding/src/encodings/physical/value.rs +++ b/rust/lance-encoding/src/encodings/physical/value.rs @@ -5,7 +5,8 @@ use arrow_buffer::{BooleanBufferBuilder, bit_util}; use crate::buffer::LanceBuffer; use crate::compression::{ - BlockCompressor, BlockDecompressor, FixedPerValueDecompressor, MiniBlockDecompressor, + BlockCompressor, BlockDecompressor, BlockValueType, FixedPerValueDecompressor, + MiniBlockDecompressor, require_block_payload, }; use crate::data::{ BlockInfo, DataBlock, FixedSizeListBlock, FixedWidthDataBlock, NullableDataBlock, @@ -13,7 +14,7 @@ use crate::data::{ use crate::encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock}; use crate::encodings::logical::primitive::miniblock::{ MAX_MINIBLOCK_BYTES, MAX_MINIBLOCK_VALUES, MiniBlockChunk, MiniBlockCompressed, - MiniBlockCompressor, + MiniBlockCompressionContext, MiniBlockCompressor, }; use crate::format::ProtobufUtils21; use crate::format::pb21::compressive_encoding::Compression; @@ -458,20 +459,66 @@ impl ValueEncoder { } impl BlockCompressor for ValueEncoder { - fn compress(&self, data: DataBlock) -> Result { + fn compress(&self, data: DataBlock) -> Result> { let data = match data { DataBlock::FixedWidth(fixed_width) => fixed_width.data, - _ => unimplemented!( - "Cannot compress block of type {} with ValueEncoder", - data.name() - ), + _ => { + return Err(Error::invalid_input(format!( + "ValueEncoder cannot compress a {} block", + data.name() + ))); + } }; - Ok(data) + Ok(Some(data)) + } +} + +/// Flat fixed-width block compressor used by the generic block selector. +#[derive(Debug)] +pub(crate) struct FixedWidthBlockCompressor { + value_type: BlockValueType, +} + +impl FixedWidthBlockCompressor { + pub(crate) fn new(value_type: BlockValueType) -> Self { + Self { value_type } + } +} + +impl BlockCompressor for FixedWidthBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + return Err(Error::invalid_input( + "Flat block compression requires fixed-width data", + )); + }; + if data.bits_per_value != self.value_type.bits_per_value() { + return Err(Error::invalid_input(format!( + "Flat block compressor expects {}-bit values, got {}", + self.value_type.bits_per_value(), + data.bits_per_value + ))); + } + let expected = usize::try_from(data.num_values) + .ok() + .and_then(|values| values.checked_mul(self.value_type.bits_per_value() as usize / 8)) + .ok_or_else(|| Error::invalid_input("Flat block payload length overflows usize"))?; + if data.data.len() != expected { + return Err(Error::invalid_input(format!( + "Flat block input has {} bytes, expected {expected}", + data.data.len() + ))); + } + Ok(Some(data.data)) } } impl MiniBlockCompressor for ValueEncoder { - fn compress(&self, chunk: DataBlock) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { + fn compress( + &self, + chunk: DataBlock, + _context: MiniBlockCompressionContext, + ) -> Result<(MiniBlockCompressed, CompressiveEncoding)> { match chunk { DataBlock::FixedWidth(fixed_width) => { let encoding = ProtobufUtils21::flat(fixed_width.bits_per_value, None); @@ -571,13 +618,48 @@ impl ValueDecompressor { } impl BlockDecompressor for ValueDecompressor { - fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Flat block")?; let block = self.buffer_to_block(data, num_values); assert_eq!(block.num_values(), num_values); Ok(block) } } +/// Flat fixed-width block decompressor with fallible payload validation. +#[derive(Debug)] +pub(crate) struct FixedWidthBlockDecompressor { + value_type: BlockValueType, +} + +impl FixedWidthBlockDecompressor { + pub(crate) fn new(value_type: BlockValueType) -> Self { + Self { value_type } + } +} + +impl BlockDecompressor for FixedWidthBlockDecompressor { + fn decompress(&self, data: Option, num_values: u64) -> Result { + let data = require_block_payload(data, "Flat block")?; + let expected = usize::try_from(num_values) + .ok() + .and_then(|values| values.checked_mul(self.value_type.bits_per_value() as usize / 8)) + .ok_or_else(|| Error::invalid_input("Flat block payload length overflows usize"))?; + if data.len() != expected { + return Err(Error::invalid_input(format!( + "Flat block payload has {} bytes, expected {expected}", + data.len() + ))); + } + Ok(DataBlock::FixedWidth(FixedWidthDataBlock { + bits_per_value: self.value_type.bits_per_value(), + num_values, + data, + block_info: BlockInfo::new(), + })) + } +} + impl MiniBlockDecompressor for ValueDecompressor { fn decompress(&self, data: Vec, num_values: u64) -> Result { let num_items = num_values * self.items_per_value; @@ -775,7 +857,9 @@ mod tests { encodings::{ logical::primitive::{ fullzip::{PerValueCompressor, PerValueDataBlock}, - miniblock::MiniBlockCompressor, + miniblock::{ + MiniBlockCompressed, MiniBlockCompressionContext, MiniBlockCompressor, + }, }, physical::value::ValueDecompressor, }, @@ -789,6 +873,16 @@ mod tests { use super::ValueEncoder; + fn compress_miniblock( + compressor: &dyn MiniBlockCompressor, + data: DataBlock, + ) -> lance_core::Result<( + MiniBlockCompressed, + crate::format::pb21::CompressiveEncoding, + )> { + compressor.compress(data, MiniBlockCompressionContext::new(0, true, true)) + } + const PRIMITIVE_TYPES: &[DataType] = &[ DataType::Null, DataType::FixedSizeBinary(2), @@ -969,7 +1063,7 @@ mod tests { let starting_data = DataBlock::from_array(sample_list.clone()); let encoder = ValueEncoder::default(); - let (data, compression) = MiniBlockCompressor::compress(&encoder, starting_data).unwrap(); + let (data, compression) = compress_miniblock(&encoder, starting_data).unwrap(); assert_eq!(data.num_values, 3); assert_eq!(data.data.len(), 3); @@ -1030,7 +1124,7 @@ mod tests { let starting_data = DataBlock::from_array(array); let encoder = ValueEncoder::default(); - let result = MiniBlockCompressor::compress(&encoder, starting_data); + let result = compress_miniblock(&encoder, starting_data); let err = result.expect_err("wide values should not be encodable as miniblock"); assert!( @@ -1142,7 +1236,7 @@ mod tests { ); let encoder = ValueEncoder::default(); - let (data, compression) = MiniBlockCompressor::compress(&encoder, starting_data).unwrap(); + let (data, compression) = compress_miniblock(&encoder, starting_data).unwrap(); let Compression::FixedSizeList(fsl) = compression.compression.unwrap() else { panic!() diff --git a/rust/lance-encoding/src/format.rs b/rust/lance-encoding/src/format.rs index f37e69b0216..f0c2ca575c6 100644 --- a/rust/lance-encoding/src/format.rs +++ b/rust/lance-encoding/src/format.rs @@ -670,6 +670,60 @@ 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 dictionary( + indices: crate::format::pb21::CompressiveEncoding, + items: crate::format::pb21::CompressiveEncoding, + num_dictionary_items: u32, + ) -> crate::format::pb21::CompressiveEncoding { + crate::format::pb21::CompressiveEncoding { + compression: Some( + crate::format::pb21::compressive_encoding::Compression::Dictionary(Box::new( + crate::format::pb21::Dictionary { + indices: Some(Box::new(indices)), + items: Some(Box::new(items)), + num_dictionary_items, + }, + )), + ), + } + } + 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() { diff --git a/rust/lance-encoding/tests/compression_strategy.rs b/rust/lance-encoding/tests/compression_strategy.rs new file mode 100644 index 00000000000..23e3a70dddf --- /dev/null +++ b/rust/lance-encoding/tests/compression_strategy.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +use arrow_schema::{DataType, Field as ArrowField}; +use lance_core::{Result, datatypes::Field}; +use lance_encoding::{ + buffer::LanceBuffer, + compression::{BlockCompressor, CompressionStrategy, DefaultCompressionStrategy}, + data::{BlockInfo, DataBlock, FixedWidthDataBlock}, + encodings::logical::primitive::{fullzip::PerValueCompressor, miniblock::MiniBlockCompressor}, + format::ProtobufUtils21, +}; + +#[derive(Debug)] +struct IdentityBlockCompressor; + +impl BlockCompressor for IdentityBlockCompressor { + fn compress(&self, data: DataBlock) -> Result> { + let DataBlock::FixedWidth(data) = data else { + panic!("test compressor only accepts fixed-width data"); + }; + Ok(Some(data.data)) + } +} + +#[derive(Debug, Default)] +struct CustomCompressionStrategy { + fallback: DefaultCompressionStrategy, +} + +impl CompressionStrategy for CustomCompressionStrategy { + fn create_block_compressor( + &self, + _field: &Field, + data: &DataBlock, + ) -> Result<( + Box, + lance_encoding::format::pb21::CompressiveEncoding, + )> { + let DataBlock::FixedWidth(data) = data else { + panic!("test strategy only accepts fixed-width data"); + }; + Ok(( + Box::new(IdentityBlockCompressor), + ProtobufUtils21::flat(data.bits_per_value, None), + )) + } + + fn create_per_value( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + self.fallback.create_per_value(field, data) + } + + fn create_miniblock_compressor( + &self, + field: &Field, + data: &DataBlock, + ) -> Result> { + self.fallback.create_miniblock_compressor(field, data) + } +} + +#[test] +fn public_strategy_returns_the_frozen_block_compressor() { + let field = Field::try_from(&ArrowField::new("values", DataType::UInt32, false)).unwrap(); + let values = vec![3_u32, 5, 8, 13]; + let data = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(values.clone()), + bits_per_value: 32, + num_values: values.len() as u64, + block_info: BlockInfo::default(), + }); + let strategy: Box = Box::new(CustomCompressionStrategy::default()); + + let (compressor, _) = strategy.create_block_compressor(&field, &data).unwrap(); + let payload = compressor.compress(data).unwrap().unwrap(); + assert_eq!(payload.borrow_to_typed_slice::().as_ref(), values); +} From afc2a51b7a33bcfecd2f269b39dad077675f225d Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 27 Jul 2026 16:11:49 +0800 Subject: [PATCH 2/2] fix(encoding): add missing license header --- rust/lance-encoding/src/compression/block/tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/rust/lance-encoding/src/compression/block/tests.rs b/rust/lance-encoding/src/compression/block/tests.rs index 230f9660ee8..16853eac926 100644 --- a/rust/lance-encoding/src/compression/block/tests.rs +++ b/rust/lance-encoding/src/compression/block/tests.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + use super::*; use std::sync::Arc;