liquid-cache 0.1.12

//! LiquidArray is the core data structure of LiquidCache.
//! You should not use this module directly.
//! Instead, use `liquid_cache_datafusion_server` or `liquid_cache_datafusion_client` to interact with LiquidCache.
mod byte_array;
pub mod byte_view_array;
mod decimal_array;
mod fix_len_byte_array;
mod float_array;
mod hybrid_primitive_array;
pub mod ipc;
mod linear_integer_array;
mod primitive_array;
pub mod raw;
mod squeezed_date32_array;
#[cfg(test)]
mod tests;
pub(crate) mod utils;
mod variant_array;

use std::{any::Any, ops::Range, sync::Arc};

use arrow::{
    array::{ArrayRef, BooleanArray},
    buffer::BooleanBuffer,
};
use arrow_schema::DataType;
pub use byte_array::{LiquidByteArray, get_bytes_needle, get_string_needle};
pub use byte_view_array::LiquidByteViewArray;
use bytes::Bytes;
use datafusion::logical_expr::Operator as DFOperator;
use datafusion::physical_plan::PhysicalExpr;
pub use decimal_array::LiquidDecimalArray;
pub use fix_len_byte_array::LiquidFixedLenByteArray;
use float_array::LiquidFloatType;
pub use float_array::{LiquidFloat32Array, LiquidFloat64Array, LiquidFloatArray};
pub use linear_integer_array::{
    LiquidLinearArray, LiquidLinearDate32Array, LiquidLinearDate64Array, LiquidLinearI8Array,
    LiquidLinearI16Array, LiquidLinearI32Array, LiquidLinearI64Array, LiquidLinearU8Array,
    LiquidLinearU16Array, LiquidLinearU32Array, LiquidLinearU64Array,
};
pub use primitive_array::IntegerSqueezePolicy;
pub use primitive_array::{
    LiquidDate32Array, LiquidDate64Array, LiquidI8Array, LiquidI16Array, LiquidI32Array,
    LiquidI64Array, LiquidPrimitiveArray, LiquidPrimitiveDeltaArray, LiquidPrimitiveType,
    LiquidU8Array, LiquidU16Array, LiquidU32Array, LiquidU64Array,
};
pub use squeezed_date32_array::{Date32Field, SqueezedDate32Array};
pub use variant_array::VariantStructSqueezedArray;

use crate::{cache::CacheExpression, liquid_array::raw::FsstArray};

/// Liquid data type is only logical type
#[derive(Debug, Clone, Copy)]
#[repr(u16)]
pub enum LiquidDataType {
    /// A byte array.
    ByteArray = 0,
    /// A byte-view array (dictionary + FSST raw + views).
    ByteViewArray = 4,
    /// An integer.
    Integer = 1,
    /// A float.
    Float = 2,
    /// A fixed length byte array.
    FixedLenByteArray = 3,
    /// A decimal encoded as a primitive u64 array.
    Decimal = 6,
    /// A linear-model based integer (signed residuals + model params).
    LinearInteger = 5,
}

impl From<u16> for LiquidDataType {
    fn from(value: u16) -> Self {
        match value {
            0 => LiquidDataType::ByteArray,
            4 => LiquidDataType::ByteViewArray,
            1 => LiquidDataType::Integer,
            2 => LiquidDataType::Float,
            3 => LiquidDataType::FixedLenByteArray,
            5 => LiquidDataType::LinearInteger,
            6 => LiquidDataType::Decimal,
            _ => panic!("Invalid liquid data type: {value}"),
        }
    }
}

/// A trait to access the underlying Liquid array.
pub trait AsLiquidArray {
    /// Get the underlying string array.
    fn as_string_array_opt(&self) -> Option<&LiquidByteArray>;

    /// Get the underlying string array.
    fn as_string(&self) -> &LiquidByteArray {
        self.as_string_array_opt().expect("liquid string array")
    }

    /// Get the underlying binary array.
    fn as_binary_array_opt(&self) -> Option<&LiquidByteArray>;

    /// Get the underlying binary array.
    fn as_binary(&self) -> &LiquidByteArray {
        self.as_binary_array_opt().expect("liquid binary array")
    }

    /// Get the underlying byte view array.
    fn as_byte_view_array_opt(&self) -> Option<&LiquidByteViewArray<FsstArray>>;

    /// Get the underlying byte view array.
    fn as_byte_view(&self) -> &LiquidByteViewArray<FsstArray> {
        self.as_byte_view_array_opt()
            .expect("liquid byte view array")
    }

    /// Get the underlying primitive array.
    fn as_primitive_array_opt<T: LiquidPrimitiveType>(&self) -> Option<&LiquidPrimitiveArray<T>>;

    /// Get the underlying primitive array.
    fn as_primitive<T: LiquidPrimitiveType>(&self) -> &LiquidPrimitiveArray<T> {
        self.as_primitive_array_opt()
            .expect("liquid primitive array")
    }

    /// Get the underlying float array.
    fn as_float_array_opt<T: LiquidFloatType>(&self) -> Option<&LiquidFloatArray<T>>;

    /// Get the underlying float array.
    fn as_float<T: LiquidFloatType>(&self) -> &LiquidFloatArray<T> {
        self.as_float_array_opt().expect("liquid float array")
    }
}

impl AsLiquidArray for dyn LiquidArray + '_ {
    fn as_string_array_opt(&self) -> Option<&LiquidByteArray> {
        self.as_any().downcast_ref()
    }

    fn as_primitive_array_opt<T: LiquidPrimitiveType>(&self) -> Option<&LiquidPrimitiveArray<T>> {
        self.as_any().downcast_ref()
    }

    fn as_binary_array_opt(&self) -> Option<&LiquidByteArray> {
        self.as_any().downcast_ref()
    }

    fn as_byte_view_array_opt(&self) -> Option<&LiquidByteViewArray<FsstArray>> {
        self.as_any().downcast_ref()
    }

    fn as_float_array_opt<T: LiquidFloatType>(&self) -> Option<&LiquidFloatArray<T>> {
        self.as_any().downcast_ref()
    }
}

/// A Liquid array.
pub trait LiquidArray: std::fmt::Debug + Send + Sync {
    /// Get the underlying any type.
    fn as_any(&self) -> &dyn Any;

    /// Get the memory size of the Liquid array.
    fn get_array_memory_size(&self) -> usize;

    /// Get the length of the Liquid array.
    fn len(&self) -> usize;

    /// Check if the Liquid array is empty.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Convert the Liquid array to an Arrow array.
    fn to_arrow_array(&self) -> ArrayRef;

    /// Convert the Liquid array to an Arrow array.
    /// Except that it will pick the best encoding for the arrow array.
    /// Meaning that it may not obey the data type of the original arrow array.
    fn to_best_arrow_array(&self) -> ArrayRef {
        self.to_arrow_array()
    }

    /// Get the logical data type of the Liquid array.
    fn data_type(&self) -> LiquidDataType;

    /// Get the original arrow data type of the Liquid array.
    fn original_arrow_data_type(&self) -> DataType;

    /// Serialize the Liquid array to a byte array.
    fn to_bytes(&self) -> Vec<u8>;

    /// Filter the Liquid array with a boolean array and return an **arrow array**.
    fn filter(&self, selection: &BooleanBuffer) -> ArrayRef {
        let arrow_array = self.to_arrow_array();
        let selection = BooleanArray::new(selection.clone(), None);
        arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap()
    }

    /// Try to evaluate a predicate on the Liquid array with a filter.
    /// Returns `None` if the predicate is not supported.
    ///
    /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable.
    /// The returned boolean mask is nullable if the the original array is nullable.
    fn try_eval_predicate(
        &self,
        _predicate: &Arc<dyn PhysicalExpr>,
        _filter: &BooleanBuffer,
    ) -> Option<BooleanArray> {
        None
    }

    /// Squeeze the Liquid array to a `LiquidHybridArrayRef` and a `bytes::Bytes`.
    /// Return `None` if the Liquid array cannot be squeezed.
    ///
    /// This is the bridge from in-memory array to hybrid array.
    /// Important: The returned `Bytes` is the data that is stored on disk, it is the same as to_bytes().
    ///
    /// Hydrating the hybrid array from the stored bytes should yield the same `LiquidArray`.
    fn squeeze(
        &self,
        _io: Arc<dyn SqueezeIoHandler>,
        _expression_hint: Option<&CacheExpression>,
    ) -> Option<(LiquidSqueezedArrayRef, bytes::Bytes)> {
        None
    }
}

/// A reference to a Liquid array.
pub type LiquidArrayRef = Arc<dyn LiquidArray>;

/// On-disk backing type for a hybrid array.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SqueezedBacking {
    /// Bytes are stored using the Liquid IPC format.
    Liquid,
    /// Bytes are stored using Arrow IPC (or another Arrow-compatible encoding).
    Arrow,
}

/// A reference to a Liquid squeezed array.
pub type LiquidSqueezedArrayRef = Arc<dyn LiquidSqueezedArray>;

/// Signals that the squeezed representation needs to be hydrated from disk.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct NeedsBacking;

/// Result type for squeezed operations that may require disk hydration.
pub type SqueezeResult<T> = Result<T, NeedsBacking>;

enum Operator {
    Eq,
    NotEq,
    Lt,
    LtEq,
    Gt,
    GtEq,
}

impl Operator {
    fn from_datafusion(op: &DFOperator) -> Option<Self> {
        let op = match op {
            DFOperator::Eq => Operator::Eq,
            DFOperator::NotEq => Operator::NotEq,
            DFOperator::Lt => Operator::Lt,
            DFOperator::LtEq => Operator::LtEq,
            DFOperator::Gt => Operator::Gt,
            DFOperator::GtEq => Operator::GtEq,
            _ => return None,
        };
        Some(op)
    }
}

/// A Liquid squeezed array is a Liquid array that part of its data is stored on disk.
/// `LiquidSqueezedArray` is more complex than in-memory `LiquidArray` because it needs to handle IO.
#[async_trait::async_trait]
pub trait LiquidSqueezedArray: std::fmt::Debug + Send + Sync {
    /// Get the underlying any type.
    fn as_any(&self) -> &dyn Any;

    /// Get the memory size of the Liquid array.
    fn get_array_memory_size(&self) -> usize;

    /// Get the length of the Liquid array.
    fn len(&self) -> usize;

    /// Check if the Liquid array is empty.
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Convert the Liquid array to an Arrow array.
    async fn to_arrow_array(&self) -> ArrayRef;

    /// Convert the Liquid array to an Arrow array.
    /// Except that it will pick the best encoding for the arrow array.
    /// Meaning that it may not obey the data type of the original arrow array.
    async fn to_best_arrow_array(&self) -> ArrayRef {
        self.to_arrow_array().await
    }

    /// Get the logical data type of the Liquid array.
    fn data_type(&self) -> LiquidDataType;

    /// Get the original arrow data type of the Liquid squeezed array.
    fn original_arrow_data_type(&self) -> DataType;

    /// Filter the Liquid array with a boolean array and return an **arrow array**.
    async fn filter(&self, selection: &BooleanBuffer) -> ArrayRef {
        let arrow_array = self.to_arrow_array().await;
        let selection = BooleanArray::new(selection.clone(), None);
        arrow::compute::kernels::filter::filter(&arrow_array, &selection).unwrap()
    }

    /// Try to evaluate a predicate on the Liquid array with a filter.
    /// Returns `Ok(None)` if the predicate is not supported.
    ///
    /// Note that the filter is a boolean buffer, not a boolean array, i.e., filter can't be nullable.
    /// The returned boolean mask is nullable if the the original array is nullable.
    async fn try_eval_predicate(
        &self,
        _predicate: &Arc<dyn PhysicalExpr>,
        _filter: &BooleanBuffer,
    ) -> Option<BooleanArray> {
        None
    }

    /// Describe how the squeezed array persists its backing bytes on disk.
    fn disk_backing(&self) -> SqueezedBacking {
        SqueezedBacking::Liquid
    }
}

/// A trait to read the backing bytes of a squeezed array from disk.
#[async_trait::async_trait]
pub trait SqueezeIoHandler: std::fmt::Debug + Send + Sync {
    /// Read the backing bytes of a squeezed array from disk.
    async fn read(&self, range: Option<Range<u64>>) -> std::io::Result<Bytes>;

    /// Trace the number of decompressions performed.
    // TODO: this is ugly.
    fn tracing_decompress_count(&self, _decompress_cnt: usize, _total_cnt: usize) {
        // Do nothing by default
    }

    /// Trace the number of IO saved by squeezing.
    // TODO: this is ugly.
    fn trace_io_saved(&self) {
        // Do nothing by default
    }
}

/// Compile-time info about primitive kind (signed vs unsigned) and bounds.
/// Implemented for all Liquid-supported primitive integer and date types.
pub trait PrimitiveKind {
    /// Whether the logical type is unsigned (true for u8/u16/u32/u64).
    const IS_UNSIGNED: bool;
    /// Maximum representable value as u64 for unsigned types (unused for signed).
    const MAX_U64: u64;
    /// Minimum representable value as i64 for signed/date types (unused for unsigned).
    const MIN_I64: i64;
    /// Maximum representable value as i64 for signed/date types (unused for unsigned).
    const MAX_I64: i64;
}

macro_rules! impl_unsigned_kind {
    ($t:ty, $max:expr) => {
        impl PrimitiveKind for $t {
            const IS_UNSIGNED: bool = true;
            const MAX_U64: u64 = $max as u64;
            const MIN_I64: i64 = 0; // unused
            const MAX_I64: i64 = 0; // unused
        }
    };
}

macro_rules! impl_signed_kind {
    ($t:ty, $min:expr, $max:expr) => {
        impl PrimitiveKind for $t {
            const IS_UNSIGNED: bool = false;
            const MAX_U64: u64 = 0; // unused
            const MIN_I64: i64 = $min as i64;
            const MAX_I64: i64 = $max as i64;
        }
    };
}

use arrow::datatypes::{
    Date32Type, Date64Type, Int8Type, Int16Type, Int32Type, Int64Type, TimestampMicrosecondType,
    TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type,
    UInt32Type, UInt64Type,
};

impl_unsigned_kind!(UInt8Type, u8::MAX);
impl_unsigned_kind!(UInt16Type, u16::MAX);
impl_unsigned_kind!(UInt32Type, u32::MAX);
impl_unsigned_kind!(UInt64Type, u64::MAX);

impl_signed_kind!(Int8Type, i8::MIN, i8::MAX);
impl_signed_kind!(Int16Type, i16::MIN, i16::MAX);
impl_signed_kind!(Int32Type, i32::MIN, i32::MAX);
impl_signed_kind!(Int64Type, i64::MIN, i64::MAX);

// Dates are logically signed in Arrow (Date32: i32 days, Date64: i64 ms)
impl_signed_kind!(Date32Type, i32::MIN, i32::MAX);
impl_signed_kind!(Date64Type, i64::MIN, i64::MAX);
impl_signed_kind!(TimestampSecondType, i64::MIN, i64::MAX);
impl_signed_kind!(TimestampMillisecondType, i64::MIN, i64::MAX);
impl_signed_kind!(TimestampMicrosecondType, i64::MIN, i64::MAX);
impl_signed_kind!(TimestampNanosecondType, i64::MIN, i64::MAX);