lance_encoding/encodings/physical/

block.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! Encodings based on traditional block compression schemes
//!
//! Traditional compressors take in a buffer and return a smaller buffer.  All encoding
//! description is shoved into the compressed buffer and the entire buffer is needed to
//! decompress any of the data.
//!
//! These encodings are not transparent, which limits our ability to use them.  In addition
//! they are often quite expensive in CPU terms.
//!
//! However, they are effective and useful for some cases.  For example, when working with large
//! variable length values (e.g. source code files) they can be very effective.
//!
//! The module introduces the `[BufferCompressor]` trait which describes the interface for a
//! traditional block compressor.  It is implemented for the most common compression schemes
//! (zstd, lz4, etc).
//!
//! There is not yet a mini-block variant of this compressor (but could easily be one) and the
//! full zip variant works by applying compression on a per-value basis (which allows it to be
//! transparent).

use arrow_buffer::ArrowNativeType;
use lance_core::{Error, Result};

use std::str::FromStr;

use crate::compression::{BlockCompressor, BlockDecompressor, require_block_payload};
use crate::encodings::physical::binary::{BinaryBlockDecompressor, VariableEncoder};
use crate::format::{
    EncodingUtils,
    metadata::{self, CompressiveEncoding},
};
use crate::{
    buffer::LanceBuffer,
    compression::VariablePerValueDecompressor,
    data::{BlockInfo, DataBlock, VariableWidthBlock},
    encodings::logical::primitive::fullzip::{PerValueCompressor, PerValueDataBlock},
};

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct CompressionConfig {
    pub(crate) scheme: CompressionScheme,
    pub(crate) level: Option<i32>,
}

impl CompressionConfig {
    /// Create a compression configuration for an encoding mechanism.
    pub fn new(scheme: CompressionScheme, level: Option<i32>) -> Self {
        Self { scheme, level }
    }

    /// Return the selected compression scheme.
    pub fn scheme(&self) -> CompressionScheme {
        self.scheme
    }

    /// Return the optional compression level.
    pub fn level(&self) -> Option<i32> {
        self.level
    }
}

impl Default for CompressionConfig {
    fn default() -> Self {
        Self {
            scheme: CompressionScheme::Lz4,
            level: Some(0),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CompressionScheme {
    None,
    Fsst,
    Zstd,
    Lz4,
}

impl TryFrom<CompressionScheme> for metadata::CompressionScheme {
    type Error = Error;

    fn try_from(scheme: CompressionScheme) -> Result<Self> {
        match scheme {
            CompressionScheme::Lz4 => Ok(Self::CompressionAlgorithmLz4),
            CompressionScheme::Zstd => Ok(Self::CompressionAlgorithmZstd),
            _ => Err(Error::invalid_input(format!(
                "Unsupported compression scheme: {:?}",
                scheme
            ))),
        }
    }
}

impl TryFrom<metadata::CompressionScheme> for CompressionScheme {
    type Error = Error;

    fn try_from(scheme: metadata::CompressionScheme) -> Result<Self> {
        match scheme {
            metadata::CompressionScheme::CompressionAlgorithmLz4 => Ok(Self::Lz4),
            metadata::CompressionScheme::CompressionAlgorithmZstd => Ok(Self::Zstd),
            _ => Err(Error::invalid_input(format!(
                "Unsupported compression scheme: {:?}",
                scheme
            ))),
        }
    }
}

impl std::fmt::Display for CompressionScheme {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        let scheme_str = match self {
            Self::Fsst => "fsst",
            Self::Zstd => "zstd",
            Self::None => "none",
            Self::Lz4 => "lz4",
        };
        write!(f, "{}", scheme_str)
    }
}

impl FromStr for CompressionScheme {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self> {
        match s {
            "none" => Ok(Self::None),
            "fsst" => Ok(Self::Fsst),
            "zstd" => Ok(Self::Zstd),
            "lz4" => Ok(Self::Lz4),
            _ => Err(Error::invalid_input(format!(
                "Unknown compression scheme: {}",
                s
            ))),
        }
    }
}

pub trait BufferCompressor: std::fmt::Debug + Send + Sync {
    fn compress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()>;
    fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()>;

    fn decompress_exact(
        &self,
        input_buf: &[u8],
        output_buf: &mut Vec<u8>,
        expected_bytes: usize,
    ) -> Result<()> {
        self.decompress(input_buf, output_buf)?;
        if output_buf.len() != expected_bytes {
            return Err(Error::invalid_input(format!(
                "General compression produced {} bytes, expected {expected_bytes}",
                output_buf.len()
            )));
        }
        Ok(())
    }

    fn config(&self) -> CompressionConfig;
}

#[cfg(feature = "zstd")]
mod zstd {
    use std::io::{Cursor, Write};
    use std::sync::{Mutex, OnceLock};

    use super::*;

    use ::zstd::bulk::{Compressor, Decompressor};
    use ::zstd::stream::copy_decode;
    use std::cell::RefCell;

    thread_local! {
        /// One reusable decompression context per thread: decode tasks run on a
        /// shared pool, so a per-thread context avoids both per-value context
        /// setup and cross-thread lock contention.
        static DECOMPRESSOR: RefCell<Option<Decompressor<'static>>> = const { RefCell::new(None) };
    }

    /// Appends into `buf`'s spare capacity (reserved by the caller) without
    /// zero-filling it first; zstd writes the bytes and the cursor sets the length.
    fn append_cursor(buf: &mut Vec<u8>) -> Cursor<&mut Vec<u8>> {
        let start = buf.len() as u64;
        let mut cursor = Cursor::new(buf);
        cursor.set_position(start);
        cursor
    }

    /// Decompresses one zstd frame into `output_buf`'s spare capacity, which the
    /// caller reserves, and returns the number of bytes appended. A frame larger
    /// than the spare capacity is an error.
    fn decompress_frame_into(
        compressed: &[u8],
        output_buf: &mut Vec<u8>,
    ) -> std::io::Result<usize> {
        DECOMPRESSOR.with(|slot| {
            let mut slot = slot.borrow_mut();
            if slot.is_none() {
                *slot = Some(Decompressor::new()?);
            }
            slot.as_mut()
                .expect("initialized above")
                .decompress_to_buffer(compressed, &mut append_cursor(output_buf))
        })
    }

    /// A zstd buffer compressor that lazily creates and reuses compression contexts.
    ///
    /// The compression context is cached to enable reuse across chunks within a
    /// page. It is lazily initialized to prevent it from getting initialized on
    /// decode-only codepaths.
    ///
    /// Decompression does not share this compressor's mutex-guarded context:
    /// decode tasks run concurrently on a shared pool, so a shared context would
    /// serialize them. Each thread instead keeps its own decompression context
    /// (see `DECOMPRESSOR`), which avoids both per-value context setup and lock
    /// contention.
    pub struct ZstdBufferCompressor {
        compression_level: i32,
        compressor: OnceLock<std::result::Result<Mutex<Compressor<'static>>, String>>,
    }

    impl std::fmt::Debug for ZstdBufferCompressor {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("ZstdBufferCompressor")
                .field("compression_level", &self.compression_level)
                .finish()
        }
    }

    impl ZstdBufferCompressor {
        pub fn new(compression_level: i32) -> Self {
            Self {
                compression_level,
                compressor: OnceLock::new(),
            }
        }

        fn get_compressor(&self) -> Result<&Mutex<Compressor<'static>>> {
            self.compressor
                .get_or_init(|| {
                    Compressor::new(self.compression_level)
                        .map(Mutex::new)
                        .map_err(|e| e.to_string())
                })
                .as_ref()
                .map_err(|e| Error::internal(format!("Failed to create zstd compressor: {}", e)))
        }

        // https://datatracker.ietf.org/doc/html/rfc8878
        fn is_raw_stream_format(&self, input_buf: &[u8]) -> bool {
            if input_buf.len() < 8 {
                return true; // can't be length prefixed format if less than 8 bytes
            }
            // read the first 4 bytes as the magic number
            let mut magic_buf = [0u8; 4];
            magic_buf.copy_from_slice(&input_buf[..4]);
            let magic = u32::from_le_bytes(magic_buf);

            // see RFC 8878, section 3.1.1. Zstandard Frames, which defines the magic number
            const ZSTD_MAGIC_NUMBER: u32 = 0xFD2FB528;
            if magic == ZSTD_MAGIC_NUMBER {
                // the compressed buffer starts like a Zstd frame.
                // Per RFC 8878, the reserved bit (with Bit Number 3, the 4th bit) in the FHD (frame header descriptor) MUST be 0
                // see section 3.1.1.1.1. 'Frame_Header_Descriptor' and section 3.1.1.1.1.4. 'Reserved Bit' for details
                const FHD_BYTE_INDEX: usize = 4;
                let fhd_byte = input_buf[FHD_BYTE_INDEX];
                const FHD_RESERVED_BIT_MASK: u8 = 0b0001_0000;
                let reserved_bit = fhd_byte & FHD_RESERVED_BIT_MASK;

                if reserved_bit != 0 {
                    // this bit is 1. This is NOT a valid zstd frame.
                    // therefore, it must be length prefixed format where the length coincidentally
                    // started with the magic number
                    false
                } else {
                    // the reserved bit is 0. This is consistent with a valid Zstd frame.
                    // treat it as raw stream format
                    true
                }
            } else {
                // doesn't start with the magic number, so it can't be the raw stream format
                false
            }
        }

        fn decompress_length_prefixed_zstd(
            &self,
            input_buf: &[u8],
            output_buf: &mut Vec<u8>,
        ) -> Result<()> {
            const LENGTH_PREFIX_SIZE: usize = 8;
            let mut len_buf = [0u8; LENGTH_PREFIX_SIZE];
            len_buf.copy_from_slice(&input_buf[..LENGTH_PREFIX_SIZE]);

            let uncompressed_len = u64::from_le_bytes(len_buf) as usize;

            let start = output_buf.len();
            output_buf.reserve(uncompressed_len);
            let compressed_data = &input_buf[LENGTH_PREFIX_SIZE..];
            let written = decompress_frame_into(compressed_data, output_buf)?;
            if written != uncompressed_len {
                // `reserve` may over-allocate, so a frame longer than its prefix can
                // still fit; drop whatever it appended so an error leaves the
                // buffer exactly as the caller passed it.
                output_buf.truncate(start);
                return Err(Error::internal(format!(
                    "Zstd decompressed {written} bytes, expected {uncompressed_len}"
                )));
            }
            Ok(())
        }
    }

    impl BufferCompressor for ZstdBufferCompressor {
        fn compress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
            output_buf.write_all(&(input_buf.len() as u64).to_le_bytes())?;

            output_buf.reserve(::zstd::zstd_safe::compress_bound(input_buf.len()));
            self.get_compressor()?
                .lock()
                .unwrap()
                .compress_to_buffer(input_buf, &mut append_cursor(output_buf))
                .map_err(|e| Error::internal(format!("Zstd compression error: {}", e)))?;
            Ok(())
        }

        fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
            if input_buf.is_empty() {
                return Ok(());
            }

            let is_raw_stream_format = self.is_raw_stream_format(input_buf);
            if is_raw_stream_format {
                copy_decode(Cursor::new(input_buf), output_buf)?;
            } else {
                self.decompress_length_prefixed_zstd(input_buf, output_buf)?;
            }

            Ok(())
        }

        fn decompress_exact(
            &self,
            input_buf: &[u8],
            output_buf: &mut Vec<u8>,
            expected_bytes: usize,
        ) -> Result<()> {
            if input_buf.is_empty() {
                if expected_bytes == 0 {
                    return Ok(());
                }
                return Err(Error::invalid_input(format!(
                    "Zstd payload is empty, expected {expected_bytes} output bytes"
                )));
            }
            let compressed = if self.is_raw_stream_format(input_buf) {
                input_buf
            } else {
                let declared = u64::from_le_bytes(input_buf[..8].try_into().map_err(|_| {
                    Error::invalid_input("Zstd payload is shorter than its 8-byte length prefix")
                })?);
                let declared = usize::try_from(declared).map_err(|_| {
                    Error::invalid_input("Zstd declared output length does not fit usize")
                })?;
                if declared != expected_bytes {
                    return Err(Error::invalid_input(format!(
                        "Zstd declares {declared} output bytes, expected {expected_bytes}"
                    )));
                }
                &input_buf[8..]
            };
            let start = output_buf.len();
            // Reports a length overflow as well as an allocation failure.
            output_buf
                .try_reserve_exact(expected_bytes)
                .map_err(|error| {
                    Error::invalid_input(format!(
                        "Zstd could not reserve {expected_bytes} output bytes: {error}"
                    ))
                })?;
            let decoded = decompress_frame_into(compressed, output_buf);
            let decoded = match decoded {
                Ok(decoded) => decoded,
                Err(error) => {
                    output_buf.truncate(start);
                    return Err(error.into());
                }
            };
            if decoded != expected_bytes {
                output_buf.truncate(start);
                return Err(Error::invalid_input(format!(
                    "Zstd produced {decoded} bytes, expected {expected_bytes}"
                )));
            }
            Ok(())
        }

        fn config(&self) -> CompressionConfig {
            CompressionConfig {
                scheme: CompressionScheme::Zstd,
                level: Some(self.compression_level),
            }
        }
    }
}

#[cfg(feature = "lz4")]
mod lz4 {
    use super::*;

    #[derive(Debug, Default)]
    pub struct Lz4BufferCompressor {}

    impl BufferCompressor for Lz4BufferCompressor {
        fn compress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
            // Remember the starting position
            let start_pos = output_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)
            output_buf.resize(start_pos + max_size + 4, 0);

            let compressed_size = ::lz4::block::compress_to_buffer(
                input_buf,
                None,
                true,
                &mut output_buf[start_pos..],
            )
            .map_err(|err| Error::internal(format!("LZ4 compression error: {}", err)))?;

            // Truncate to actual size
            output_buf.truncate(start_pos + compressed_size);
            Ok(())
        }

        fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
            // When prepend_size is true, LZ4 stores the uncompressed size in the first 4 bytes
            // We can read this to know exactly how much space we need
            if input_buf.len() < 4 {
                return Err(Error::invalid_input(
                    "LZ4 buffer is missing its length prefix",
                ));
            }

            // Read the uncompressed size from the first 4 bytes (little-endian)
            let uncompressed_size =
                u32::from_le_bytes([input_buf[0], input_buf[1], input_buf[2], input_buf[3]])
                    as usize;

            // Remember the starting position
            let start_pos = output_buf.len();

            // Resize to ensure we have the exact space needed
            output_buf.resize(start_pos + uncompressed_size, 0);

            // Now decompress directly into the buffer slice
            let decompressed_size =
                ::lz4::block::decompress_to_buffer(input_buf, None, &mut output_buf[start_pos..])
                    .map_err(|err| Error::internal(format!("LZ4 decompression error: {}", err)))?;

            // Truncate to actual decompressed size (should be same as uncompressed_size)
            output_buf.truncate(start_pos + decompressed_size);

            Ok(())
        }

        fn decompress_exact(
            &self,
            input_buf: &[u8],
            output_buf: &mut Vec<u8>,
            expected_bytes: usize,
        ) -> Result<()> {
            let declared = input_buf.get(..4).ok_or_else(|| {
                Error::invalid_input("LZ4 payload is shorter than its 4-byte length prefix")
            })?;
            let declared =
                u32::from_le_bytes(declared.try_into().map_err(|_| {
                    Error::invalid_input("LZ4 payload has an invalid length prefix")
                })?) as usize;
            if declared != expected_bytes {
                return Err(Error::invalid_input(format!(
                    "LZ4 declares {declared} output bytes, expected {expected_bytes}"
                )));
            }
            let start = output_buf.len();
            let end = start
                .checked_add(expected_bytes)
                .ok_or_else(|| Error::invalid_input("LZ4 output buffer length overflows usize"))?;
            output_buf
                .try_reserve_exact(expected_bytes)
                .map_err(|error| {
                    Error::invalid_input(format!(
                        "LZ4 could not reserve {expected_bytes} output bytes: {error}"
                    ))
                })?;
            output_buf.resize(end, 0);
            let decoded =
                ::lz4::block::decompress_to_buffer(input_buf, None, &mut output_buf[start..]);
            let decoded = match decoded {
                Ok(decoded) => decoded,
                Err(error) => {
                    output_buf.truncate(start);
                    return Err(Error::invalid_input(format!(
                        "LZ4 decompression error: {error}"
                    )));
                }
            };
            if decoded != expected_bytes {
                output_buf.truncate(start);
                return Err(Error::invalid_input(format!(
                    "LZ4 produced {decoded} bytes, expected {expected_bytes}"
                )));
            }
            Ok(())
        }

        fn config(&self) -> CompressionConfig {
            CompressionConfig {
                scheme: CompressionScheme::Lz4,
                level: None,
            }
        }
    }
}

#[derive(Debug, Default)]
pub struct NoopBufferCompressor {}

impl BufferCompressor for NoopBufferCompressor {
    fn compress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
        output_buf.extend_from_slice(input_buf);
        Ok(())
    }

    fn decompress(&self, input_buf: &[u8], output_buf: &mut Vec<u8>) -> Result<()> {
        output_buf.extend_from_slice(input_buf);
        Ok(())
    }

    fn decompress_exact(
        &self,
        input_buf: &[u8],
        output_buf: &mut Vec<u8>,
        expected_bytes: usize,
    ) -> Result<()> {
        if input_buf.len() != expected_bytes {
            return Err(Error::invalid_input(format!(
                "Uncompressed payload has {} bytes, expected {expected_bytes}",
                input_buf.len()
            )));
        }
        output_buf.extend_from_slice(input_buf);
        Ok(())
    }

    fn config(&self) -> CompressionConfig {
        CompressionConfig {
            scheme: CompressionScheme::None,
            level: None,
        }
    }
}

pub struct GeneralBufferCompressor {}

impl GeneralBufferCompressor {
    pub fn get_compressor(
        compression_config: CompressionConfig,
    ) -> Result<Box<dyn BufferCompressor>> {
        match compression_config.scheme {
            // FSST has its own compression path and isn't implemented as a generic buffer compressor
            CompressionScheme::Fsst => Err(Error::invalid_input_source(
                "fsst is not usable as a general buffer compressor".into(),
            )),
            CompressionScheme::Zstd => {
                #[cfg(feature = "zstd")]
                {
                    Ok(Box::new(zstd::ZstdBufferCompressor::new(
                        compression_config.level.unwrap_or(0),
                    )))
                }
                #[cfg(not(feature = "zstd"))]
                {
                    Err(Error::invalid_input_source(
                        "package was not built with zstd support".into(),
                    ))
                }
            }
            CompressionScheme::Lz4 => {
                #[cfg(feature = "lz4")]
                {
                    Ok(Box::new(lz4::Lz4BufferCompressor::default()))
                }
                #[cfg(not(feature = "lz4"))]
                {
                    Err(Error::invalid_input_source(
                        "package was not built with lz4 support".into(),
                    ))
                }
            }
            CompressionScheme::None => Ok(Box::new(NoopBufferCompressor {})),
        }
    }
}

/// A block decompressor that first applies general-purpose compression (LZ4/Zstd)
/// before delegating to an inner block decompressor.
#[derive(Debug)]
pub struct GeneralBlockDecompressor {
    inner: Box<dyn BlockDecompressor>,
    compressor: Box<dyn BufferCompressor>,
}

impl GeneralBlockDecompressor {
    pub fn try_new(
        inner: Box<dyn BlockDecompressor>,
        compression: CompressionConfig,
    ) -> Result<Self> {
        let compressor = GeneralBufferCompressor::get_compressor(compression)?;
        Ok(Self { inner, compressor })
    }
}

impl BlockDecompressor for GeneralBlockDecompressor {
    fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
        let data = require_block_payload(data, "General block compression")?;
        let mut decompressed = Vec::new();
        self.compressor.decompress(&data, &mut decompressed)?;
        self.inner
            .decompress(Some(LanceBuffer::from(decompressed)), num_values)
    }
}

// An encoder which uses generic compression, such as zstd/lz4 to encode buffers
#[derive(Debug)]
pub struct CompressedBufferEncoder {
    pub(crate) compressor: Box<dyn BufferCompressor>,
    // Runtime compressors normalize levels that they default or ignore. Block descriptors must
    // retain the selected configuration so stable writers preserve those present/absent values.
    block_compression: CompressionConfig,
}

impl Default for CompressedBufferEncoder {
    fn default() -> Self {
        // Pick zstd if available, otherwise lz4, otherwise none
        #[cfg(feature = "zstd")]
        let (scheme, level) = (CompressionScheme::Zstd, Some(0));
        #[cfg(all(feature = "lz4", not(feature = "zstd")))]
        let (scheme, level) = (CompressionScheme::Lz4, None);
        #[cfg(not(any(feature = "zstd", feature = "lz4")))]
        let (scheme, level) = (CompressionScheme::None, None);

        let block_compression = CompressionConfig { scheme, level };
        let compressor = GeneralBufferCompressor::get_compressor(block_compression).unwrap();
        Self {
            compressor,
            block_compression,
        }
    }
}

impl CompressedBufferEncoder {
    pub fn try_new(compression_config: CompressionConfig) -> Result<Self> {
        let compressor = GeneralBufferCompressor::get_compressor(compression_config)?;
        Ok(Self {
            compressor,
            block_compression: compression_config,
        })
    }

    pub fn from_scheme(scheme: metadata::CompressionScheme) -> Result<Self> {
        let scheme = CompressionScheme::try_from(scheme)?;
        let block_compression = CompressionConfig {
            scheme,
            level: Some(0),
        };
        Ok(Self {
            compressor: GeneralBufferCompressor::get_compressor(block_compression)?,
            block_compression,
        })
    }
}

impl CompressedBufferEncoder {
    pub fn per_value_compress<T: ArrowNativeType>(
        &self,
        data: &[u8],
        offsets: &[T],
        compressed: &mut Vec<u8>,
    ) -> Result<LanceBuffer> {
        let mut new_offsets: Vec<T> = Vec::with_capacity(offsets.len());
        new_offsets.push(T::from_usize(0).unwrap());

        for off in offsets.windows(2) {
            let start = off[0].as_usize();
            let end = off[1].as_usize();
            self.compressor.compress(&data[start..end], compressed)?;
            new_offsets.push(T::from_usize(compressed.len()).unwrap());
        }

        Ok(LanceBuffer::reinterpret_vec(new_offsets))
    }

    pub fn per_value_decompress<T: ArrowNativeType>(
        &self,
        data: &[u8],
        offsets: &[T],
        decompressed: &mut Vec<u8>,
    ) -> Result<LanceBuffer> {
        let mut new_offsets: Vec<T> = Vec::with_capacity(offsets.len());
        new_offsets.push(T::from_usize(0).unwrap());

        for off in offsets.windows(2) {
            let start = off[0].as_usize();
            let end = off[1].as_usize();
            self.compressor
                .decompress(&data[start..end], decompressed)?;
            new_offsets.push(T::from_usize(decompressed.len()).unwrap());
        }

        Ok(LanceBuffer::reinterpret_vec(new_offsets))
    }
}

impl PerValueCompressor for CompressedBufferEncoder {
    fn compress(&self, data: DataBlock) -> Result<(PerValueDataBlock, CompressiveEncoding)> {
        let data_type = data.name();
        let data = data.as_variable_width().ok_or(Error::internal(format!(
            "Attempt to use CompressedBufferEncoder on data of type {}",
            data_type
        )))?;

        let data_bytes = &data.data;
        let mut compressed = Vec::with_capacity(data_bytes.len());

        let new_offsets = match data.bits_per_offset {
            32 => self.per_value_compress::<u32>(
                data_bytes,
                &data.offsets.borrow_to_typed_slice::<u32>(),
                &mut compressed,
            )?,
            64 => self.per_value_compress::<u64>(
                data_bytes,
                &data.offsets.borrow_to_typed_slice::<u64>(),
                &mut compressed,
            )?,
            _ => unreachable!(),
        };

        let compressed = PerValueDataBlock::Variable(VariableWidthBlock {
            bits_per_offset: data.bits_per_offset,
            data: LanceBuffer::from(compressed),
            offsets: new_offsets,
            num_values: data.num_values,
            block_info: BlockInfo::new(),
        });

        // TODO: Support setting the level
        // TODO: Support underlying compression of data (e.g. defer to binary encoding for offset bitpacking)
        let encoding = EncodingUtils::wrapped(
            self.compressor.config(),
            EncodingUtils::variable(EncodingUtils::flat(data.bits_per_offset as u64, None), None),
        )?;

        Ok((compressed, encoding))
    }
}

impl VariablePerValueDecompressor for CompressedBufferEncoder {
    fn decompress(&self, data: VariableWidthBlock) -> Result<DataBlock> {
        let data_bytes = &data.data;
        let mut decompressed = Vec::with_capacity(data_bytes.len() * 2);

        let new_offsets = match data.bits_per_offset {
            32 => self.per_value_decompress(
                data_bytes,
                &data.offsets.borrow_to_typed_slice::<u32>(),
                &mut decompressed,
            )?,
            64 => self.per_value_decompress(
                data_bytes,
                &data.offsets.borrow_to_typed_slice::<u64>(),
                &mut decompressed,
            )?,
            _ => unreachable!(),
        };
        Ok(DataBlock::VariableWidth(VariableWidthBlock {
            bits_per_offset: data.bits_per_offset,
            data: LanceBuffer::from(decompressed),
            offsets: new_offsets,
            num_values: data.num_values,
            block_info: BlockInfo::new(),
        }))
    }
}

impl BlockCompressor for CompressedBufferEncoder {
    fn compress(&self, data: DataBlock) -> Result<(Option<LanceBuffer>, CompressiveEncoding)> {
        let (encoded, inner_encoding) = match data {
            DataBlock::FixedWidth(fixed_width) => (
                fixed_width.data,
                EncodingUtils::flat(fixed_width.bits_per_value, None),
            ),
            DataBlock::VariableWidth(variable_width) => {
                // Wrap VariableEncoder to handle the encoding
                let encoder = VariableEncoder::default();
                let (payload, encoding) =
                    BlockCompressor::compress(&encoder, DataBlock::VariableWidth(variable_width))?;
                (
                    payload.ok_or_else(|| {
                        Error::internal(
                            "VariableEncoder returned no payload for general compression"
                                .to_string(),
                        )
                    })?,
                    encoding,
                )
            }
            _ => {
                return Err(Error::invalid_input_source(
                    "Unsupported data block type".into(),
                ));
            }
        };

        let mut compressed = Vec::new();
        self.compressor.compress(&encoded, &mut compressed)?;
        Ok((
            Some(LanceBuffer::from(compressed)),
            EncodingUtils::wrapped(self.block_compression, inner_encoding)?,
        ))
    }
}

impl BlockDecompressor for CompressedBufferEncoder {
    fn decompress(&self, data: Option<LanceBuffer>, num_values: u64) -> Result<DataBlock> {
        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(Some(LanceBuffer::from(decompressed)), num_values)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::str::FromStr;

    use crate::encodings::physical::block::zstd::ZstdBufferCompressor;

    #[test]
    fn test_compression_scheme_from_str() {
        assert_eq!(
            CompressionScheme::from_str("none").unwrap(),
            CompressionScheme::None
        );
        assert_eq!(
            CompressionScheme::from_str("zstd").unwrap(),
            CompressionScheme::Zstd
        );
    }

    #[test]
    fn test_compression_scheme_from_str_invalid() {
        assert!(CompressionScheme::from_str("invalid").is_err());
    }

    #[cfg(feature = "zstd")]
    mod zstd {
        use std::io::Write;
        use std::sync::Arc;

        use super::*;

        #[test]
        fn test_compress_zstd_with_length_prefixed() {
            let compressor = ZstdBufferCompressor::new(0);
            let input_data = b"Hello, world!";
            let mut compressed_data = Vec::new();

            compressor
                .compress(input_data, &mut compressed_data)
                .unwrap();
            let mut decompressed_data = Vec::new();
            compressor
                .decompress(&compressed_data, &mut decompressed_data)
                .unwrap();
            assert_eq!(input_data, decompressed_data.as_slice());
        }

        /// Deterministic pseudo-random bytes (xorshift64), so inputs are reproducible
        /// without a test-only RNG dependency.
        fn pseudo_random_bytes(seed: u64, len: usize) -> Vec<u8> {
            let mut state = seed | 1;
            (0..len)
                .map(|_| {
                    state ^= state << 13;
                    state ^= state >> 7;
                    state ^= state << 17;
                    state as u8
                })
                .collect()
        }

        /// Inputs covering the shapes zstd treats differently: empty, tiny,
        /// incompressible, highly repetitive, and all-zero, up to 1 MiB.
        fn zstd_test_inputs() -> Vec<Vec<u8>> {
            let mut inputs = Vec::new();
            for (i, &size) in [0usize, 1, 7, 64, 4095, 65_536, 1 << 20].iter().enumerate() {
                inputs.push(pseudo_random_bytes(i as u64 + 1, size));
                inputs.push((0..size).map(|j| b"lance zstd "[j % 11]).collect());
                inputs.push(vec![0u8; size]);
            }
            inputs
        }

        /// A length-prefixed frame built with the `zstd` crate directly, independent
        /// of `ZstdBufferCompressor`: the on-disk layout the decompressor must read.
        fn reference_frame(input: &[u8], level: i32) -> Vec<u8> {
            let mut frame = (input.len() as u64).to_le_bytes().to_vec();
            frame.extend(::zstd::bulk::compress(input, level).unwrap());
            frame
        }

        /// Compression and decompression append into the buffer's spare capacity
        /// without zero-filling it first. Round-trip across sizes (including empty
        /// and incompressible inputs), levels, pre-existing buffer contents, and
        /// spare capacity left over from earlier writes; the prefix must survive
        /// and every byte past it must be exactly what zstd produced.
        #[test]
        fn test_zstd_append_into_spare_capacity_round_trips() {
            for level in [0, 1, 3, 9] {
                let compressor = ZstdBufferCompressor::new(level);
                for input in zstd_test_inputs() {
                    let prefix = b"existing-prefix".to_vec();
                    let mut compressed = Vec::with_capacity(prefix.len() + 3);
                    compressed.extend_from_slice(&prefix);
                    compressor.compress(&input, &mut compressed).unwrap();
                    assert_eq!(&compressed[..prefix.len()], &prefix[..]);

                    let mut decompressed = Vec::with_capacity(4);
                    decompressed.extend_from_slice(b"xy");
                    compressor
                        .decompress(&compressed[prefix.len()..], &mut decompressed)
                        .unwrap();
                    assert_eq!(&decompressed[..2], b"xy");
                    assert_eq!(&decompressed[2..], &input[..]);
                }
            }
        }

        /// The written bytes are exactly what one-shot zstd produces for the same
        /// level, behind the same 8-byte length prefix, so appending into spare
        /// capacity cannot change what lands on disk.
        #[test]
        fn test_zstd_compress_output_matches_reference_frame() {
            for level in [0, 1, 3, 9, 19] {
                let compressor = ZstdBufferCompressor::new(level);
                for input in zstd_test_inputs() {
                    let mut compressed = Vec::new();
                    compressor.compress(&input, &mut compressed).unwrap();
                    assert_eq!(
                        compressed,
                        reference_frame(&input, level),
                        "level {level}, {} input bytes",
                        input.len()
                    );
                }
            }
        }

        /// Frames written by plain zstd decode through `decompress`, and frames
        /// written by `compress` decode with plain zstd.
        #[test]
        fn test_zstd_interoperates_with_plain_zstd() {
            let compressor = ZstdBufferCompressor::new(3);
            for input in zstd_test_inputs() {
                let mut decompressed = Vec::new();
                compressor
                    .decompress(&reference_frame(&input, 3), &mut decompressed)
                    .unwrap();
                assert_eq!(decompressed, input);

                let mut compressed = Vec::new();
                compressor.compress(&input, &mut compressed).unwrap();
                let plain = ::zstd::bulk::decompress(&compressed[8..], input.len()).unwrap();
                assert_eq!(plain, input);
            }
        }

        /// A length prefix that disagrees with the frame is an error in both
        /// directions, and the output buffer is left exactly as it was passed in
        /// (including when the frame is longer and fits in over-reserved capacity).
        #[test]
        fn test_zstd_decompress_length_mismatch_errors_and_preserves_buffer() {
            let compressor = ZstdBufferCompressor::new(0);
            let input = pseudo_random_bytes(7, 10_000);
            let frame = reference_frame(&input, 3);
            for declared in [0u64, 1, 9_999, 10_001, 20_000] {
                let mut bad = frame.clone();
                bad[..8].copy_from_slice(&declared.to_le_bytes());
                for spare in [0usize, 64 * 1024] {
                    let mut output = Vec::with_capacity(4 + spare);
                    output.extend_from_slice(b"keep");
                    assert!(
                        compressor.decompress(&bad, &mut output).is_err(),
                        "declared {declared} bytes for a {}-byte frame must fail",
                        input.len()
                    );
                    assert_eq!(output, b"keep", "declared {declared}, spare {spare}");
                }
            }
        }

        /// Corrupted or truncated frames fail without exposing any bytes.
        #[test]
        fn test_zstd_decompress_corrupt_frame_errors_and_preserves_buffer() {
            let compressor = ZstdBufferCompressor::new(0);
            let input = pseudo_random_bytes(11, 50_000);
            let frame = reference_frame(&input, 3);
            let mut corrupted = frame.clone();
            corrupted[8..12].copy_from_slice(&[0, 0, 0, 0]);
            let truncated = frame[..frame.len() / 2].to_vec();
            for bad in [corrupted, truncated] {
                let mut output = b"keep".to_vec();
                assert!(compressor.decompress(&bad, &mut output).is_err());
                assert_eq!(output, b"keep");
            }
        }

        /// The per-thread decompression context is reset for every frame: a failed
        /// frame does not poison the next one, and frames of very different sizes
        /// decode correctly back to back on the same thread.
        #[test]
        fn test_zstd_decompressor_context_reuse_on_one_thread() {
            let compressor = ZstdBufferCompressor::new(0);
            let big = pseudo_random_bytes(3, 1 << 20);
            let tiny = b"x".to_vec();
            let empty = Vec::new();
            let mut corrupt = reference_frame(&big, 3);
            corrupt[8..12].copy_from_slice(&[0, 0, 0, 0]);
            for input in [&big, &tiny, &empty, &big, &tiny] {
                let mut scratch = Vec::new();
                assert!(compressor.decompress(&corrupt, &mut scratch).is_err());
                let mut output = Vec::new();
                compressor
                    .decompress(&reference_frame(input, 3), &mut output)
                    .unwrap();
                assert_eq!(&output, input);
            }
        }

        /// Many threads decompress concurrently through one shared compressor; each
        /// thread uses its own context and every frame decodes to its input.
        #[test]
        fn test_zstd_decompress_concurrently_from_many_threads() {
            let compressor = Arc::new(ZstdBufferCompressor::new(0));
            let handles: Vec<_> = (0..8u64)
                .map(|thread| {
                    let compressor = Arc::clone(&compressor);
                    std::thread::spawn(move || {
                        for i in 0..64u64 {
                            let len = ((thread * 64 + i) * 7919 % 200_000) as usize;
                            let input = pseudo_random_bytes(thread * 1000 + i, len);
                            let mut compressed = Vec::new();
                            compressor.compress(&input, &mut compressed).unwrap();
                            let mut output = Vec::new();
                            compressor.decompress(&compressed, &mut output).unwrap();
                            assert_eq!(output, input);
                        }
                    })
                })
                .collect();
            for handle in handles {
                handle.join().unwrap();
            }
        }

        /// Output buffers with no spare capacity (every append must grow the
        /// allocation) and with far more capacity than needed both produce the
        /// same bytes, with the length covering only what zstd wrote.
        #[test]
        fn test_zstd_output_buffer_capacity_edges() {
            let compressor = ZstdBufferCompressor::new(0);
            for input in zstd_test_inputs() {
                let expected = reference_frame(&input, 0);

                let mut tight = b"prefix".to_vec();
                tight.shrink_to_fit();
                compressor.compress(&input, &mut tight).unwrap();
                assert_eq!(&tight[..6], b"prefix");
                assert_eq!(&tight[6..], &expected[..]);

                let mut roomy = Vec::with_capacity(8 << 20);
                roomy.extend_from_slice(b"prefix");
                compressor.compress(&input, &mut roomy).unwrap();
                assert_eq!(roomy.len(), 6 + expected.len());
                assert_eq!(&roomy[6..], &expected[..]);

                let mut decoded = Vec::with_capacity(8 << 20);
                compressor.decompress(&expected, &mut decoded).unwrap();
                assert_eq!(decoded, input);
            }
        }

        /// Decompressing an empty input is a no-op, as before.
        #[test]
        fn test_zstd_decompress_empty_input_is_noop() {
            let compressor = ZstdBufferCompressor::new(0);
            let mut output = b"keep".to_vec();
            compressor.decompress(&[], &mut output).unwrap();
            assert_eq!(output, b"keep");
        }

        #[test]
        fn test_zstd_exact_output_rejects_declared_length() {
            let compressor = ZstdBufferCompressor::new(0);
            let mut compressed = Vec::new();
            compressor.compress(b"bounded", &mut compressed).unwrap();
            compressed[..8].copy_from_slice(&u64::MAX.to_le_bytes());

            let mut output = Vec::new();
            let error = compressor
                .decompress_exact(&compressed, &mut output, 7)
                .unwrap_err();
            assert!(error.to_string().contains("expected 7"));
            assert!(output.is_empty());
        }

        /// Exact decompression appends after existing contents, for length-prefixed
        /// and raw frames, regardless of how much spare capacity the buffer has.
        #[test]
        fn test_zstd_exact_output_appends_into_spare_capacity() {
            let compressor = ZstdBufferCompressor::new(0);
            for input in zstd_test_inputs() {
                let prefixed = reference_frame(&input, 3);
                let raw = ::zstd::stream::encode_all(input.as_slice(), 3).unwrap();
                for frame in [&prefixed, &raw] {
                    for spare in [0usize, 4 << 20] {
                        let mut output = Vec::with_capacity(4 + spare);
                        output.extend_from_slice(b"keep");
                        compressor
                            .decompress_exact(frame, &mut output, input.len())
                            .unwrap();
                        assert_eq!(&output[..4], b"keep");
                        assert_eq!(&output[4..], &input[..]);
                    }
                }
            }
        }

        /// A frame that decodes to more or fewer bytes than expected fails and
        /// leaves the buffer as it was, even when spare capacity could hold the
        /// longer frame.
        #[test]
        fn test_zstd_exact_output_frame_size_mismatch_preserves_buffer() {
            let compressor = ZstdBufferCompressor::new(0);
            let input = pseudo_random_bytes(5, 10_000);
            let raw = ::zstd::stream::encode_all(input.as_slice(), 3).unwrap();
            for expected in [0usize, 1, 9_999, 10_001] {
                let mut prefixed = reference_frame(&input, 3);
                prefixed[..8].copy_from_slice(&(expected as u64).to_le_bytes());
                for frame in [&prefixed, &raw] {
                    for spare in [0usize, 64 * 1024] {
                        let mut output = Vec::with_capacity(4 + spare);
                        output.extend_from_slice(b"keep");
                        assert!(
                            compressor
                                .decompress_exact(frame, &mut output, expected)
                                .is_err(),
                            "expected {expected} bytes from a {}-byte frame must fail",
                            input.len()
                        );
                        assert_eq!(output, b"keep", "expected {expected}, spare {spare}");
                    }
                }
            }
        }

        #[test]
        fn test_zstd_compress_decompress_multiple_times() {
            let compressor = ZstdBufferCompressor::new(0);
            let (input_data_1, input_data_2) = (b"Hello ", b"World");
            let mut compressed_data = Vec::new();

            compressor
                .compress(input_data_1, &mut compressed_data)
                .unwrap();
            let compressed_length_1 = compressed_data.len();

            compressor
                .compress(input_data_2, &mut compressed_data)
                .unwrap();

            let mut decompressed_data = Vec::new();
            compressor
                .decompress(
                    &compressed_data[..compressed_length_1],
                    &mut decompressed_data,
                )
                .unwrap();

            compressor
                .decompress(
                    &compressed_data[compressed_length_1..],
                    &mut decompressed_data,
                )
                .unwrap();

            // the output should contain both input_data_1 and input_data_2
            assert_eq!(
                decompressed_data.len(),
                input_data_1.len() + input_data_2.len()
            );
            assert_eq!(
                &decompressed_data[..input_data_1.len()],
                input_data_1,
                "First part of decompressed data should match input_1"
            );
            assert_eq!(
                &decompressed_data[input_data_1.len()..],
                input_data_2,
                "Second part of decompressed data should match input_2"
            );
        }

        #[test]
        fn test_compress_zstd_raw_stream_format_and_decompress_with_length_prefixed() {
            let compressor = ZstdBufferCompressor::new(0);
            let input_data = b"Hello, world!";
            let mut compressed_data = Vec::new();

            // compress using raw stream format
            let mut encoder = ::zstd::Encoder::new(&mut compressed_data, 0).unwrap();
            encoder.write_all(input_data).unwrap();
            encoder.finish().expect("failed to encode data with zstd");

            // decompress using length prefixed format
            let mut decompressed_data = Vec::new();
            compressor
                .decompress(&compressed_data, &mut decompressed_data)
                .unwrap();
            assert_eq!(input_data, decompressed_data.as_slice());
        }
    }

    #[cfg(feature = "lz4")]
    mod lz4 {
        use std::{collections::HashMap, sync::Arc};

        use arrow_schema::{DataType, Field};
        use lance_datagen::array::{binary_prefix_plus_counter, utf8_prefix_plus_counter};

        use super::*;

        use crate::constants::DICT_SIZE_RATIO_META_KEY;
        use crate::{
            constants::{
                COMPRESSION_META_KEY, DICT_DIVISOR_META_KEY, STRUCTURAL_ENCODING_FULLZIP,
                STRUCTURAL_ENCODING_META_KEY,
            },
            encodings::physical::block::lz4::Lz4BufferCompressor,
            testing::{
                FnArrayGeneratorProvider, TestCases, TestEncoding,
                check_round_trip_encoding_generated,
            },
        };

        #[test]
        fn test_lz4_compress_decompress() {
            let compressor = Lz4BufferCompressor::default();
            let input_data = b"Hello, world!";
            let mut compressed_data = Vec::new();

            compressor
                .compress(input_data, &mut compressed_data)
                .unwrap();
            let mut decompressed_data = Vec::new();
            compressor
                .decompress(&compressed_data, &mut decompressed_data)
                .unwrap();
            assert_eq!(input_data, decompressed_data.as_slice());
        }

        #[test]
        fn test_lz4_exact_output_rejects_declared_length() {
            let compressor = Lz4BufferCompressor::default();
            let mut compressed = Vec::new();
            compressor.compress(b"bounded", &mut compressed).unwrap();
            compressed[..4].copy_from_slice(&u32::MAX.to_le_bytes());

            let mut output = Vec::new();
            let error = compressor
                .decompress_exact(&compressed, &mut output, 7)
                .unwrap_err();
            assert!(error.to_string().contains("declares"));
            assert!(output.is_empty());
        }

        #[rstest::rstest]
        #[test_log::test(tokio::test)]
        async fn test_lz4_compress_round_trip(
            #[values(
                DataType::Utf8,
                DataType::LargeUtf8,
                DataType::Binary,
                DataType::LargeBinary
            )]
            data_type: DataType,
            #[values(
                TestEncoding::StructuralU16,
                TestEncoding::StructuralU32,
                TestEncoding::StructuralSparse
            )]
            encoding: TestEncoding,
            #[values(false, true)] use_slicing: bool,
        ) {
            let field = Field::new("", data_type.clone(), false);
            let mut field_meta = HashMap::new();
            field_meta.insert(COMPRESSION_META_KEY.to_string(), "lz4".to_string());
            // Some bad cardinality estimatation causes us to use dictionary encoding currently
            // which causes the expected encoding check to fail.
            field_meta.insert(DICT_DIVISOR_META_KEY.to_string(), "100000".to_string());
            field_meta.insert(DICT_SIZE_RATIO_META_KEY.to_string(), "0.0001".to_string());
            // Also disable size-based dictionary encoding
            field_meta.insert(
                STRUCTURAL_ENCODING_META_KEY.to_string(),
                STRUCTURAL_ENCODING_FULLZIP.to_string(),
            );
            let field = field.with_metadata(field_meta);
            let test_cases = TestCases::basic()
                // Need to use large pages as small pages might be too small to compress
                .with_page_sizes(vec![1024 * 1024])
                .with_expected_encoding("zstd")
                .with_encoding(encoding)
                .with_slicing_modes([use_slicing]);

            // Can't use the default random provider because random data isn't compressible
            // and we will fallback to uncompressed encoding
            let datagen = Box::new(FnArrayGeneratorProvider::new(move || match data_type {
                DataType::Utf8 => utf8_prefix_plus_counter("compressme", false),
                DataType::Binary => {
                    binary_prefix_plus_counter(Arc::from(b"compressme".to_owned()), false)
                }
                DataType::LargeUtf8 => utf8_prefix_plus_counter("compressme", true),
                DataType::LargeBinary => {
                    binary_prefix_plus_counter(Arc::from(b"compressme".to_owned()), true)
                }
                _ => panic!("Unsupported data type: {:?}", data_type),
            }));

            check_round_trip_encoding_generated(field, datagen, test_cases).await;
        }
    }
}