diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 15b39d62b26..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::{ @@ -106,6 +107,51 @@ 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"; +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 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 { + 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 + } + #[cfg(not(feature = "lz4"))] + { + u64::MAX + } +} struct PageLoadTask { decoder_fut: BoxFuture<'static, Result>>, @@ -5418,6 +5464,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 +5472,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 +5486,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 +5496,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 +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)?; + 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)?; @@ -8943,11 +8994,87 @@ 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 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() { + let over = super::lz4_max_input_size() + 1; + let metadata = PrimitiveStructuralEncoder::resolve_dict_values_compression_metadata( + &HashMap::new(), + None, + None, + over, + ); + let expected = if cfg!(feature = "zstd") { + "zstd" + } else { + "none" + }; + assert_eq!( + metadata.get(COMPRESSION_META_KEY), + 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. + 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()),); + } + + /// 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")] + #[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 +9091,7 @@ mod tests { &field_metadata, Some("zstd".to_string()), Some("3".to_string()), + 1024, ); assert_eq!( metadata.get(COMPRESSION_META_KEY), @@ -8981,6 +9109,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/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 { 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();