vortex_btrblocks/

lib.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

#![deny(missing_docs)]

//! Vortex's [BtrBlocks]-inspired adaptive compression framework.
//!
//! This crate provides a sophisticated multi-level compression system that adaptively selects
//! optimal compression schemes based on data characteristics. The compressor analyzes arrays
//! to determine the best encoding strategy, supporting cascaded compression with multiple
//! encoding layers for maximum efficiency.
//!
//! # Key Features
//!
//! - **Adaptive Compression**: Automatically selects the best compression scheme based on data patterns
//! - **Type-Specific Compressors**: Specialized compression for integers, floats, strings, and temporal data
//! - **Cascaded Encoding**: Multiple compression layers can be applied for optimal results
//! - **Statistical Analysis**: Uses data sampling and statistics to predict compression ratios
//! - **Recursive Structure Handling**: Compresses nested structures like structs and lists
//!
//! # Example
//!
//! ```rust
//! use vortex_btrblocks::BtrBlocksCompressor;
//! use vortex_array::Array;
//!
//! let compressor = BtrBlocksCompressor::default();
//! // let compressed = compressor.compress(&array)?;
//! ```
//!
//! [BtrBlocks]: https://www.cs.cit.tum.de/fileadmin/w00cfj/dis/papers/btrblocks.pdf

use std::fmt::Debug;
use std::hash::Hash;

use vortex_array::Array;
use vortex_array::ArrayRef;
use vortex_array::Canonical;
use vortex_array::IntoArray;
use vortex_array::ToCanonical;
use vortex_array::arrays::ConstantArray;
use vortex_array::arrays::ExtensionArray;
use vortex_array::arrays::FixedSizeListArray;
use vortex_array::arrays::ListArray;
use vortex_array::arrays::StructArray;
use vortex_array::arrays::TemporalArray;
use vortex_array::arrays::list_from_list_view;
use vortex_array::compute::Cost;
use vortex_array::vtable::VTable;
use vortex_array::vtable::ValidityHelper;
use vortex_dtype::DType;
use vortex_dtype::Nullability;
use vortex_dtype::datetime::TemporalMetadata;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;

use crate::decimal::compress_decimal;
pub use crate::float::FloatCompressor;
pub use crate::float::FloatStats;
pub use crate::float::dictionary::dictionary_encode as float_dictionary_encode;
pub use crate::integer::IntCompressor;
pub use crate::integer::IntegerStats;
pub use crate::integer::dictionary::dictionary_encode as integer_dictionary_encode;
pub use crate::string::StringCompressor;
pub use crate::string::StringStats;
pub use crate::temporal::compress_temporal;

mod decimal;
mod float;
mod integer;
mod patches;
mod rle;
mod sample;
mod string;
mod temporal;

/// Configures how stats are generated.
pub struct GenerateStatsOptions {
    /// Should distinct values should be counted during stats generation.
    pub count_distinct_values: bool,
    // pub count_runs: bool,
    // should this be scheme-specific?
}

impl Default for GenerateStatsOptions {
    fn default() -> Self {
        Self {
            count_distinct_values: true,
            // count_runs: true,
        }
    }
}

const SAMPLE_SIZE: u32 = 64;

/// Stats for the compressor.
pub trait CompressorStats: Debug + Clone {
    /// The type of the underlying source array vtable.
    type ArrayVTable: VTable;

    /// Generates stats with default options.
    fn generate(input: &<Self::ArrayVTable as VTable>::Array) -> Self {
        Self::generate_opts(input, GenerateStatsOptions::default())
    }

    /// Generates stats with provided options.
    fn generate_opts(
        input: &<Self::ArrayVTable as VTable>::Array,
        opts: GenerateStatsOptions,
    ) -> Self;

    /// Returns the underlying source array that statistics were generated from.
    fn source(&self) -> &<Self::ArrayVTable as VTable>::Array;

    /// Sample the array with default options.
    fn sample(&self, sample_size: u32, sample_count: u32) -> Self {
        self.sample_opts(sample_size, sample_count, GenerateStatsOptions::default())
    }

    /// Sample the array with provided options.
    fn sample_opts(&self, sample_size: u32, sample_count: u32, opts: GenerateStatsOptions) -> Self;
}

/// Top-level compression scheme trait.
///
/// Variants are specialized for each data type, e.g. see `IntegerScheme`, `FloatScheme`, etc.
pub trait Scheme: Debug {
    /// Type of the stats generated by the compression scheme.
    type StatsType: CompressorStats;
    /// Type of the code used to uniquely identify the compression scheme.
    type CodeType: Copy + Eq + Hash;

    /// Scheme unique identifier.
    fn code(&self) -> Self::CodeType;

    /// True if this is the singular Constant scheme for this data type.
    fn is_constant(&self) -> bool {
        false
    }

    /// Estimate the compression ratio for running this scheme (and its children)
    /// for the given input.
    ///
    /// Depth is the depth in the encoding tree we've already reached before considering this
    /// scheme.
    ///
    /// Returns the estimated compression ratio as well as the tree of compressors to use.
    fn expected_compression_ratio(
        &self,
        stats: &Self::StatsType,
        is_sample: bool,
        allowed_cascading: usize,
        excludes: &[Self::CodeType],
    ) -> VortexResult<f64> {
        estimate_compression_ratio_with_sampling(
            self,
            stats,
            is_sample,
            allowed_cascading,
            excludes,
        )
    }

    /// Compress the input with this scheme, yielding a new array.
    fn compress(
        &self,
        stats: &Self::StatsType,
        is_sample: bool,
        allowed_cascading: usize,
        excludes: &[Self::CodeType],
    ) -> VortexResult<ArrayRef>;
}

fn estimate_compression_ratio_with_sampling<T: Scheme + ?Sized>(
    compressor: &T,
    stats: &T::StatsType,
    is_sample: bool,
    allowed_cascading: usize,
    excludes: &[T::CodeType],
) -> VortexResult<f64> {
    let sample = if is_sample {
        stats.clone()
    } else {
        // We want to sample about 1% of data
        let source_len = stats.source().len();

        // We want to sample about 1% of data, while keeping a minimal sample of 640 values.
        let sample_count = usize::max(
            (source_len / 100)
                / usize::try_from(SAMPLE_SIZE).vortex_expect("SAMPLE_SIZE must fit in usize"),
            10,
        );

        tracing::trace!(
            "Sampling {} values out of {}",
            SAMPLE_SIZE as usize * sample_count,
            source_len
        );

        stats.sample(
            SAMPLE_SIZE,
            sample_count
                .try_into()
                .vortex_expect("sample count must fit in u32"),
        )
    };

    let after = compressor
        .compress(&sample, true, allowed_cascading, excludes)?
        .nbytes();
    let before = sample.source().nbytes();

    tracing::debug!(
        "estimate_compression_ratio_with_sampling(compressor={compressor:#?} is_sample={is_sample}, allowed_cascading={allowed_cascading}) = {}",
        before as f64 / after as f64
    );

    Ok(before as f64 / after as f64)
}

const MAX_CASCADE: usize = 3;

/// A compressor for a particular input type.
///
/// This trait defines the interface for type-specific compressors that can adaptively
/// choose and apply compression schemes based on data characteristics. Compressors
/// analyze input arrays, select optimal compression schemes, and handle cascading
/// compression with multiple encoding layers.
///
/// The compressor works by generating statistics on the input data, evaluating
/// available compression schemes, and selecting the one with the best compression ratio.
pub trait Compressor {
    /// The VTable type for arrays this compressor operates on.
    type ArrayVTable: VTable;
    /// The compression scheme type used by this compressor.
    type SchemeType: Scheme<StatsType = Self::StatsType> + ?Sized;
    /// The statistics type used to analyze arrays for compression.
    type StatsType: CompressorStats<ArrayVTable = Self::ArrayVTable>;

    /// Returns all available compression schemes for this compressor.
    fn schemes() -> &'static [&'static Self::SchemeType];
    /// Returns the default fallback compression scheme.
    fn default_scheme() -> &'static Self::SchemeType;
    /// Returns the scheme code for dictionary compression.
    fn dict_scheme_code() -> <Self::SchemeType as Scheme>::CodeType;

    /// Compresses an array using the optimal compression scheme.
    ///
    /// Generates statistics on the input array, selects the best compression scheme,
    /// and applies it. Returns the original array if compression would increase size.
    fn compress(
        array: &<Self::ArrayVTable as VTable>::Array,
        is_sample: bool,
        allowed_cascading: usize,
        excludes: &[<Self::SchemeType as Scheme>::CodeType],
    ) -> VortexResult<ArrayRef>
    where
        Self::SchemeType: 'static,
    {
        // Avoid compressing empty arrays.
        if array.is_empty() {
            return Ok(array.to_array());
        }

        // Generate stats on the array directly.
        let stats = if excludes.contains(&Self::dict_scheme_code()) {
            Self::StatsType::generate_opts(
                array,
                GenerateStatsOptions {
                    count_distinct_values: false,
                },
            )
        } else {
            Self::StatsType::generate(array)
        };
        let best_scheme = Self::choose_scheme(&stats, is_sample, allowed_cascading, excludes)?;

        let output = best_scheme.compress(&stats, is_sample, allowed_cascading, excludes)?;
        if output.nbytes() < array.nbytes() {
            Ok(output)
        } else {
            tracing::debug!("resulting tree too large: {}", output.display_tree());
            Ok(array.to_array())
        }
    }

    /// Selects the best compression scheme based on expected compression ratios.
    ///
    /// Evaluates all available schemes against the provided statistics and returns
    /// the one with the highest compression ratio. Falls back to the default scheme
    /// if no scheme provides compression benefits.
    #[allow(clippy::cognitive_complexity)]
    fn choose_scheme(
        stats: &Self::StatsType,
        is_sample: bool,
        allowed_cascading: usize,
        excludes: &[<Self::SchemeType as Scheme>::CodeType],
    ) -> VortexResult<&'static Self::SchemeType> {
        let mut best_ratio = 1.0;
        let mut best_scheme: Option<&'static Self::SchemeType> = None;

        // logging helpers
        let depth = MAX_CASCADE - allowed_cascading;

        for scheme in Self::schemes().iter() {
            if excludes.contains(&scheme.code()) {
                continue;
            }

            // We never choose Constant for a sample
            if is_sample && scheme.is_constant() {
                continue;
            }

            tracing::trace!(is_sample, depth, ?scheme, "Trying compression scheme");

            let ratio =
                scheme.expected_compression_ratio(stats, is_sample, allowed_cascading, excludes)?;
            tracing::trace!(
                is_sample,
                depth,
                ratio,
                ?scheme,
                "Expected compression result"
            );

            if !(ratio.is_subnormal() || ratio.is_infinite() || ratio.is_nan()) {
                if ratio > best_ratio {
                    best_ratio = ratio;
                    best_scheme = Some(*scheme);
                }
            } else {
                tracing::trace!(
                    "Calculated invalid compression ratio {ratio} for scheme: {scheme:?}. Must not be sub-normal, infinite or nan."
                );
            }
        }

        tracing::trace!(depth, scheme = ?best_scheme, ratio = best_ratio, "best scheme found");

        if let Some(best) = best_scheme {
            Ok(best)
        } else {
            Ok(Self::default_scheme())
        }
    }
}

/// The main compressor type implementing BtrBlocks-inspired compression.
///
/// This compressor applies adaptive compression schemes to arrays based on their data types
/// and characteristics. It recursively compresses nested structures like structs and lists,
/// and chooses optimal compression schemes for primitive types.
///
/// The compressor works by:
/// 1. Canonicalizing input arrays to a standard representation
/// 2. Analyzing data characteristics to choose optimal compression schemes
/// 3. Recursively compressing nested structures
/// 4. Applying type-specific compression for primitives, strings, and temporal data
///
/// # Examples
///
/// ```rust
/// use vortex_btrblocks::BtrBlocksCompressor;
/// use vortex_array::Array;
///
/// let compressor = BtrBlocksCompressor::default();
/// // let compressed = compressor.compress(&array)?;
/// ```
#[derive(Default, Debug, Clone)]
pub struct BtrBlocksCompressor {
    /// Whether to exclude ints from dictionary encoding.
    ///
    /// When `true`, integer arrays will not use dictionary compression schemes,
    /// which can be useful when the data has high cardinality or when dictionary
    /// overhead would exceed compression benefits.
    pub exclude_int_dict_encoding: bool,
}

impl BtrBlocksCompressor {
    /// Compresses an array using BtrBlocks-inspired compression.
    ///
    /// First canonicalizes and compacts the array, then applies optimal compression schemes.
    pub fn compress(&self, array: &dyn Array) -> VortexResult<ArrayRef> {
        // Canonicalize the array
        let canonical = array.to_canonical();

        // Compact it, removing any wasted space before we attempt to compress it
        let compact = canonical.compact()?;

        self.compress_canonical(compact)
    }

    /// Compresses a canonical array by dispatching to type-specific compressors.
    ///
    /// Recursively compresses nested structures and applies optimal schemes for each data type.
    pub fn compress_canonical(&self, array: Canonical) -> VortexResult<ArrayRef> {
        match array {
            Canonical::Null(null_array) => Ok(null_array.into_array()),
            // TODO(aduffy): Sparse, other bool compressors.
            Canonical::Bool(bool_array) => Ok(bool_array.into_array()),
            Canonical::Primitive(primitive) => {
                if primitive.ptype().is_int() {
                    if self.exclude_int_dict_encoding {
                        IntCompressor::compress_no_dict(&primitive, false, MAX_CASCADE, &[])
                    } else {
                        IntCompressor::compress(&primitive, false, MAX_CASCADE, &[])
                    }
                } else {
                    FloatCompressor::compress(&primitive, false, MAX_CASCADE, &[])
                }
            }
            Canonical::Decimal(decimal) => compress_decimal(&decimal),
            Canonical::Struct(struct_array) => {
                let fields = struct_array
                    .fields()
                    .iter()
                    .map(|field| self.compress(field))
                    .collect::<Result<Vec<_>, _>>()?;

                Ok(StructArray::try_new(
                    struct_array.names().clone(),
                    fields,
                    struct_array.len(),
                    struct_array.validity().clone(),
                )?
                .into_array())
            }
            Canonical::List(list_view_array) => {
                // TODO(joe): We might want to write list views in the future and chose between
                // list and list view.
                let list_array = list_from_list_view(list_view_array);

                // Reset the offsets to remove garbage data that might prevent us from narrowing our
                // offsets (there could be a large amount of trailing garbage data that the current
                // views do not reference at all).
                let list_array = list_array.reset_offsets(true)?;

                let compressed_elems = self.compress(list_array.elements())?;

                // Note that since the type of our offsets are not encoded in our `DType`, and since
                // we guarantee above that all elements are referenced by offsets, we may narrow the
                // widths.

                let compressed_offsets = IntCompressor::compress_no_dict(
                    &list_array.offsets().to_primitive().narrow()?,
                    false,
                    MAX_CASCADE,
                    &[],
                )?;

                Ok(ListArray::try_new(
                    compressed_elems,
                    compressed_offsets,
                    list_array.validity().clone(),
                )?
                .into_array())
            }
            Canonical::FixedSizeList(fsl_array) => {
                let compressed_elems = self.compress(fsl_array.elements())?;

                Ok(FixedSizeListArray::try_new(
                    compressed_elems,
                    fsl_array.list_size(),
                    fsl_array.validity().clone(),
                    fsl_array.len(),
                )?
                .into_array())
            }
            Canonical::VarBinView(strings) => {
                if strings
                    .dtype()
                    .eq_ignore_nullability(&DType::Utf8(Nullability::NonNullable))
                {
                    StringCompressor::compress(&strings, false, MAX_CASCADE, &[])
                } else {
                    // Binary arrays do not compress
                    Ok(strings.into_array())
                }
            }
            Canonical::Extension(ext_array) => {
                // We compress Timestamp-level arrays with DateTimeParts compression
                if let Ok(temporal_array) = TemporalArray::try_from(ext_array.to_array())
                    && let TemporalMetadata::Timestamp(..) = temporal_array.temporal_metadata()
                {
                    if temporal_array.as_ref().is_constant_opts(Cost::Canonicalize) {
                        return Ok(ConstantArray::new(
                            temporal_array.as_ref().scalar_at(0),
                            ext_array.len(),
                        )
                        .into_array());
                    }
                    return compress_temporal(temporal_array);
                }

                // Compress the underlying storage array.
                let compressed_storage = self.compress(ext_array.storage())?;

                Ok(
                    ExtensionArray::new(ext_array.ext_dtype().clone(), compressed_storage)
                        .into_array(),
                )
            }
        }
    }
}