From b8db7f12c39e8ed7449fe207d873d0d140d521cf Mon Sep 17 00:00:00 2001 From: Beinan Date: Thu, 6 Aug 2026 20:37:20 +0000 Subject: [PATCH 1/2] fix: avoid LZ4 overflow when compressing large dictionary values Dictionary values are compressed with a single LZ4 call, but the default compression for that buffer was hardcoded to lz4 with no size check. Once the dictionary exceeds LZ4_MAX_INPUT_SIZE (0x7E000000) the write fails with a bare "Compression input too long" from the lz4 crate. This is reachable from ordinary writes: a column whose values are duplicate-heavy enough to be dictionary encoded builds a dictionary bounded only by 0.8 * data_size, so a page over ~2.6 GiB can produce an oversized dictionary buffer. Reproduced with a list column of low-entropy values, which fails on a 2.25 GB dictionary. Fall back to zstd for the default when the dictionary values buffer is too large for LZ4. Zstd has no comparable input limit and compresses these buffers better in practice. An explicitly requested scheme is still honored so a user's choice is never silently changed. Also validate the input length in Lz4BufferCompressor::compress. The lz4 crate's error names neither the actual size nor the limit, which makes the failure hard to diagnose; the error now reports both and points at zstd. Note LZ4 buffers here prepend the uncompressed length as a u32, so this format cannot represent larger inputs regardless of compress_bound. Co-Authored-By: Claude Opus 5 --- .../src/encodings/logical/primitive.rs | 83 ++++++++++++++++++- .../src/encodings/physical/block.rs | 40 +++++++++ 2 files changed, 120 insertions(+), 3 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 15b39d62b26..ae94dcd4744 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -106,6 +106,32 @@ const DEFAULT_DICT_DIVISOR: u64 = 2; const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000; const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8; const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4"; +const DEFAULT_LARGE_DICT_VALUES_COMPRESSION: &str = "zstd"; + +/// Pick the default compression for a dictionary values buffer. +/// +/// The buffer is compressed with a single call, so LZ4 cannot be used once it +/// exceeds `LZ4_MAX_INPUT_SIZE`. Zstd has no such limit and compresses this kind +/// of buffer better anyway, so fall back to it rather than failing the write. +fn default_dict_values_compression(dict_values_size: u64) -> &'static str { + if dict_values_size > lz4_max_input_size() { + DEFAULT_LARGE_DICT_VALUES_COMPRESSION + } else { + DEFAULT_DICT_VALUES_COMPRESSION + } +} + +const fn lz4_max_input_size() -> u64 { + #[cfg(feature = "lz4")] + { + crate::encodings::physical::block::LZ4_MAX_INPUT_SIZE as u64 + } + // Without the lz4 feature the default below is never selected anyway. + #[cfg(not(feature = "lz4"))] + { + u64::MAX + } +} struct PageLoadTask { decoder_fut: BoxFuture<'static, Result>>, @@ -5418,6 +5444,7 @@ impl PrimitiveStructuralEncoder { field_metadata: &HashMap, env_compression: Option, env_compression_level: Option, + dict_values_size: u64, ) -> HashMap { let mut metadata = HashMap::new(); @@ -5425,7 +5452,7 @@ impl PrimitiveStructuralEncoder { .get(DICT_VALUES_COMPRESSION_META_KEY) .cloned() .or(env_compression) - .unwrap_or_else(|| DEFAULT_DICT_VALUES_COMPRESSION.to_string()); + .unwrap_or_else(|| default_dict_values_compression(dict_values_size).to_string()); metadata.insert(COMPRESSION_META_KEY.to_string(), compression); if let Some(compression_level) = field_metadata @@ -5439,7 +5466,7 @@ impl PrimitiveStructuralEncoder { metadata } - fn build_dict_values_compressor_field(field: &Field) -> Result { + fn build_dict_values_compressor_field(field: &Field, dict_values_size: u64) -> Result { // This is an internal synthetic field used only to feed metadata into // `create_block_compressor` for dictionary values. The concrete type/name here // are not semantically meaningful; we rely on explicit metadata below to control @@ -5449,6 +5476,7 @@ impl PrimitiveStructuralEncoder { &field.metadata, env::var(DICT_VALUES_COMPRESSION_ENV_VAR).ok(), env::var(DICT_VALUES_COMPRESSION_LEVEL_ENV_VAR).ok(), + dict_values_size, ); Ok(dict_values_field) } @@ -5538,7 +5566,8 @@ impl PrimitiveStructuralEncoder { if let Some(dictionary_data) = dictionary_data { let num_dictionary_items = dictionary_data.num_values(); - let dict_values_field = Self::build_dict_values_compressor_field(field)?; + let dict_values_field = + Self::build_dict_values_compressor_field(field, dictionary_data.data_size())?; let (compressor, dictionary_encoding) = compression_strategy .create_block_compressor(&dict_values_field, &dictionary_data)?; @@ -8943,11 +8972,57 @@ mod tests { &HashMap::new(), None, None, + 1024, ); assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),); assert!(!metadata.contains_key(COMPRESSION_LEVEL_META_KEY)); } + /// Dictionary values are compressed with a single LZ4 call, which fails above + /// `LZ4_MAX_INPUT_SIZE`. Oversized buffers must fall back to zstd instead. + #[cfg(feature = "lz4")] + #[test] + fn test_resolve_dict_values_compression_metadata_large_falls_back_to_zstd() { + let over = super::lz4_max_input_size() + 1; + let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata( + &HashMap::new(), + None, + None, + over, + ); + assert_eq!( + metadata.get(COMPRESSION_META_KEY), + Some(&"zstd".to_string()), + ); + + // Exactly at the limit LZ4 is still valid, so the default is unchanged. + let at_limit = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata( + &HashMap::new(), + None, + None, + super::lz4_max_input_size(), + ); + assert_eq!(at_limit.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),); + } + + /// An explicit request must win even when the buffer is too large for LZ4, so + /// that the failure is reported rather than silently changing the user's choice. + #[cfg(feature = "lz4")] + #[test] + fn test_resolve_dict_values_compression_metadata_large_respects_explicit_request() { + let field_metadata = HashMap::from([( + DICT_VALUES_COMPRESSION_META_KEY.to_string(), + "lz4".to_string(), + )]); + let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata( + &field_metadata, + None, + None, + super::lz4_max_input_size() + 1, + ); + assert_eq!(metadata.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),); + } + #[test] fn test_resolve_dict_values_compression_metadata_metadata_overrides_env() { let field_metadata = HashMap::from([ @@ -8964,6 +9039,7 @@ mod tests { &field_metadata, Some("zstd".to_string()), Some("3".to_string()), + 1024, ); assert_eq!( metadata.get(COMPRESSION_META_KEY), @@ -8981,6 +9057,7 @@ mod tests { &HashMap::new(), Some("zstd".to_string()), Some("9".to_string()), + 1024, ); assert_eq!( metadata.get(COMPRESSION_META_KEY), diff --git a/rust/lance-encoding/src/encodings/physical/block.rs b/rust/lance-encoding/src/encodings/physical/block.rs index 188af50c8a0..ed748062a6f 100644 --- a/rust/lance-encoding/src/encodings/physical/block.rs +++ b/rust/lance-encoding/src/encodings/physical/block.rs @@ -299,10 +299,33 @@ mod zstd { } } +#[cfg(feature = "lz4")] +pub use lz4::LZ4_MAX_INPUT_SIZE; + #[cfg(feature = "lz4")] mod lz4 { use super::*; + /// The largest input a single LZ4 block compression call accepts + /// (`LZ4_MAX_INPUT_SIZE`). Beyond this `compress_bound` fails. The buffers we + /// write also prepend the uncompressed length as a `u32`, so this format cannot + /// represent larger inputs either way. + pub const LZ4_MAX_INPUT_SIZE: usize = 0x7E000000; + + /// Reject oversized input up front. The underlying crate would otherwise fail + /// with a bare "Compression input too long" that names neither the actual size + /// nor the limit, which is hard to act on. + pub fn check_input_size(len: usize) -> Result<()> { + if len > LZ4_MAX_INPUT_SIZE { + return Err(Error::invalid_input(format!( + "LZ4 compression input is {} bytes which exceeds the maximum LZ4 input \ + size of {} bytes. Use zstd for buffers of this size.", + len, LZ4_MAX_INPUT_SIZE, + ))); + } + Ok(()) + } + #[derive(Debug, Default)] pub struct Lz4BufferCompressor {} @@ -311,6 +334,8 @@ mod lz4 { // Remember the starting position let start_pos = output_buf.len(); + check_input_size(input_buf.len())?; + // LZ4 needs space for the compressed data let max_size = ::lz4::block::compress_bound(input_buf.len())?; // Resize to ensure we have enough space (including 4 bytes for size header) @@ -778,6 +803,21 @@ mod tests { testing::{FnArrayGeneratorProvider, TestCases, check_round_trip_encoding_generated}, }; + /// Guard the single-call LZ4 size limit. Checked as a pure function so the + /// test does not need to allocate multiple GiB. + #[test] + fn test_lz4_input_size_limit() { + use crate::encodings::physical::block::lz4::check_input_size; + + assert!(check_input_size(LZ4_MAX_INPUT_SIZE).is_ok()); + let err = check_input_size(LZ4_MAX_INPUT_SIZE + 1).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains(&LZ4_MAX_INPUT_SIZE.to_string()) && msg.contains("zstd"), + "error should name the limit and the alternative, got: {msg}" + ); + } + #[test] fn test_lz4_compress_decompress() { let compressor = Lz4BufferCompressor::default(); From 8d8dafd2a06cf0a785fea04805fdcb6fe211eef2 Mon Sep 17 00:00:00 2001 From: Beinan Date: Thu, 6 Aug 2026 21:54:02 +0000 Subject: [PATCH 2/2] fix: use exact serialized size and available codecs for dict values Address review feedback on the dictionary values compression fallback. The size check compared LZ4's input limit against DataBlock::data_size(), but CompressedBufferEncoder compresses the serialized form, which is larger by the header VariableEncoder prepends (16 bytes for 64-bit offsets). A dictionary just under the limit would still overflow once serialized. Add variable_encoded_size() next to the encoder it mirrors and check against that instead. The fallback also named zstd unconditionally, which fails with "package was not built with zstd support" on builds without that feature. Fall back to uncompressed storage in that case so the write still succeeds. Co-Authored-By: Claude Opus 5 --- .../src/encodings/logical/primitive.rs | 74 ++++++++++++++++--- .../src/encodings/physical/binary.rs | 11 +++ 2 files changed, 74 insertions(+), 11 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index ae94dcd4744..05173130dca 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -40,6 +40,7 @@ use lance_core::{ use log::{debug, trace}; use crate::encodings::logical::primitive::miniblock::MiniBlockChunk; +use crate::encodings::physical::binary::variable_encoded_size; use crate::encodings::physical::rle::{RleDecompressor, RleRuns}; use crate::utils::bytepack::ByteUnpacker; use crate::{ @@ -107,26 +108,45 @@ const DEFAULT_DICT_MAX_CARDINALITY: u64 = 100_000; const DEFAULT_DICT_SIZE_RATIO: f64 = 0.8; const DEFAULT_DICT_VALUES_COMPRESSION: &str = "lz4"; const DEFAULT_LARGE_DICT_VALUES_COMPRESSION: &str = "zstd"; +const UNCOMPRESSED_DICT_VALUES_COMPRESSION: &str = "none"; /// Pick the default compression for a dictionary values buffer. /// -/// The buffer is compressed with a single call, so LZ4 cannot be used once it -/// exceeds `LZ4_MAX_INPUT_SIZE`. Zstd has no such limit and compresses this kind -/// of buffer better anyway, so fall back to it rather than failing the write. -fn default_dict_values_compression(dict_values_size: u64) -> &'static str { - if dict_values_size > lz4_max_input_size() { +/// The buffer is compressed with a single call, so LZ4 cannot be used once the +/// serialized block exceeds `LZ4_MAX_INPUT_SIZE`. Prefer zstd in that case: it has +/// no comparable input limit and compresses these buffers better in practice. If +/// this build has no zstd support, fall back to storing the values uncompressed +/// rather than failing the write. +fn default_dict_values_compression(serialized_size: u64) -> &'static str { + if serialized_size <= lz4_max_input_size() { + return DEFAULT_DICT_VALUES_COMPRESSION; + } + if cfg!(feature = "zstd") { DEFAULT_LARGE_DICT_VALUES_COMPRESSION } else { - DEFAULT_DICT_VALUES_COMPRESSION + UNCOMPRESSED_DICT_VALUES_COMPRESSION } } +/// The number of bytes `block` occupies once serialized for general compression. +/// +/// `CompressedBufferEncoder` compresses this serialized form, not the raw block, so +/// codec input limits must be checked against this rather than `data_size()`. +fn dict_values_serialized_size(block: &DataBlock) -> u64 { + match block { + DataBlock::VariableWidth(variable_width) => variable_encoded_size(variable_width), + // Fixed-width blocks are compressed as-is. + other => other.data_size(), + } +} + +/// The largest serialized block a single LZ4 call accepts, or `u64::MAX` when this +/// build has no LZ4 support (in which case the LZ4 default is never selected). const fn lz4_max_input_size() -> u64 { #[cfg(feature = "lz4")] { crate::encodings::physical::block::LZ4_MAX_INPUT_SIZE as u64 } - // Without the lz4 feature the default below is never selected anyway. #[cfg(not(feature = "lz4"))] { u64::MAX @@ -5566,8 +5586,10 @@ impl PrimitiveStructuralEncoder { if let Some(dictionary_data) = dictionary_data { let num_dictionary_items = dictionary_data.num_values(); - let dict_values_field = - Self::build_dict_values_compressor_field(field, dictionary_data.data_size())?; + let dict_values_field = Self::build_dict_values_compressor_field( + field, + dict_values_serialized_size(&dictionary_data), + )?; let (compressor, dictionary_encoding) = compression_strategy .create_block_compressor(&dict_values_field, &dictionary_data)?; @@ -8979,7 +9001,8 @@ mod tests { } /// Dictionary values are compressed with a single LZ4 call, which fails above - /// `LZ4_MAX_INPUT_SIZE`. Oversized buffers must fall back to zstd instead. + /// `LZ4_MAX_INPUT_SIZE`. Oversized buffers must fall back to a codec that can + /// handle them instead of failing the write. #[cfg(feature = "lz4")] #[test] fn test_resolve_dict_values_compression_metadata_large_falls_back_to_zstd() { @@ -8990,9 +9013,15 @@ mod tests { None, over, ); + let expected = if cfg!(feature = "zstd") { + "zstd" + } else { + "none" + }; assert_eq!( metadata.get(COMPRESSION_META_KEY), - Some(&"zstd".to_string()), + Some(&expected.to_string()), + "oversized dictionary values must not default to a codec that cannot encode them" ); // Exactly at the limit LZ4 is still valid, so the default is unchanged. @@ -9005,6 +9034,29 @@ mod tests { assert_eq!(at_limit.get(COMPRESSION_META_KEY), Some(&"lz4".to_string()),); } + /// The LZ4 limit applies to the serialized block, which carries a header beyond + /// `data_size()`. A block just under the limit still serializes past it and must + /// not be handed to LZ4. + #[cfg(feature = "lz4")] + #[test] + fn test_dict_values_serialized_size_accounts_for_header() { + use crate::data::VariableWidthBlock; + + // Two 64-bit offsets, so the serialized form adds 16 header bytes. + let block = DataBlock::VariableWidth(VariableWidthBlock { + bits_per_offset: 64, + data: LanceBuffer::empty(), + offsets: LanceBuffer::reinterpret_vec(vec![0u64, 0u64]), + num_values: 1, + block_info: BlockInfo::new(), + }); + assert_eq!( + super::dict_values_serialized_size(&block), + block.data_size() + 16, + "serialized size must include the header VariableEncoder prepends" + ); + } + /// An explicit request must win even when the buffer is too large for LZ4, so /// that the failure is reported rather than silently changing the user's choice. #[cfg(feature = "lz4")] diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index 06df60434e9..d717976f475 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -452,6 +452,17 @@ impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { #[derive(Debug, Default)] pub struct VariableEncoder {} +/// The exact size of the buffer [`VariableEncoder`] produces for `block`. +/// +/// Callers that must respect a codec input limit need the serialized length, which +/// is larger than [`VariableWidthBlock::data_size`] by the header this encoder +/// prepends. Keep in sync with the `compress` implementation below. +pub fn variable_encoded_size(block: &VariableWidthBlock) -> u64 { + // bits-per-offset and bytes-start-offset, one word each. + let header_bytes = 2 * (block.bits_per_offset as u64 / 8); + header_bytes + block.offsets.len() as u64 + block.data.len() as u64 +} + impl BlockCompressor for VariableEncoder { fn compress(&self, mut data: DataBlock) -> Result { match data {