fix: avoid LZ4 overflow when compressing large dictionary values

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<Box<dyn StructuralPageDecoder>>>,
@@ -5418,14 +5464,15 @@ impl PrimitiveStructuralEncoder {
         field_metadata: &HashMap<String, String>,
         env_compression: Option<String>,
         env_compression_level: Option<String>,
+        dict_values_size: u64,
     ) -> HashMap<String, String> {
         let mut metadata = HashMap::new();
 
         let compression = field_metadata
             .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<Field> {
+    fn build_dict_values_compressor_field(field: &Field, dict_values_size: u64) -> Result<Field> {
         // 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),

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<LanceBuffer> {
         match data {

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();