liquid_cache/liquid_array/byte_view_array/

helpers.rs

use arrow::array::{
    BooleanArray, BooleanBufferBuilder, UInt16Array, cast::AsArray, types::UInt16Type,
};
use arrow::buffer::BooleanBuffer;
use datafusion_physical_expr::PhysicalExpr;
use std::sync::Arc;

use super::LiquidByteViewArray;
use super::operator::{ByteViewExpression, ByteViewOperator};
use crate::liquid_array::byte_view_array::operator::UnsupportedExpression;
use crate::liquid_array::raw::FsstArray;
use crate::liquid_array::raw::fsst_buffer::{DiskBuffer, FsstBacking};

pub(super) fn build_dict_selection(
    keys: &UInt16Array,
    dict_len: usize,
) -> (Vec<usize>, UInt16Array) {
    let mut hit_mask = BooleanBufferBuilder::new(dict_len);
    hit_mask.advance(dict_len);
    for value in keys.iter().flatten() {
        hit_mask.set_bit(value as usize, true);
    }
    let hit_mask = hit_mask.finish();
    let selected_cnt = hit_mask.count_set_bits();

    let mut key_map = vec![u16::MAX; dict_len];
    let mut selected = Vec::with_capacity(selected_cnt);
    let mut remapped: u16 = 0;
    for (index, selected_flag) in hit_mask.iter().enumerate() {
        if selected_flag {
            key_map[index] = remapped;
            selected.push(index);
            remapped += 1;
        }
    }

    let new_keys = UInt16Array::from_iter(
        keys.iter()
            .map(|value| value.map(|value| key_map[value as usize])),
    );
    (selected, new_keys)
}

pub(super) fn filter_inner<B: FsstBacking>(
    array: &LiquidByteViewArray<B>,
    filter: &BooleanBuffer,
) -> LiquidByteViewArray<B> {
    // Only filter the dictionary keys, not the offsets!
    // Offset views reference unique values in FSST buffer and should remain unchanged

    // Filter the dictionary keys using Arrow's built-in filter functionality
    let filter = BooleanArray::new(filter.clone(), None);
    let filtered_keys = arrow::compute::filter(&array.dictionary_keys, &filter).unwrap();
    let filtered_keys = filtered_keys.as_primitive::<UInt16Type>().clone();

    LiquidByteViewArray {
        dictionary_keys: filtered_keys,
        prefix_keys: array.prefix_keys.clone(),
        fsst_buffer: array.fsst_buffer.clone(),
        original_arrow_type: array.original_arrow_type,
        shared_prefix: array.shared_prefix.clone(),
        string_fingerprints: array.string_fingerprints.clone(),
    }
}

pub(super) fn try_eval_predicate_in_memory(
    expr: &Arc<dyn PhysicalExpr>,
    array: &LiquidByteViewArray<FsstArray>,
) -> Option<BooleanArray> {
    let expr = match ByteViewExpression::try_from(expr) {
        Ok(expr) => expr,
        Err(UnsupportedExpression::Constant(v)) => {
            let bool_array = if v {
                BooleanBuffer::new_set(array.len())
            } else {
                BooleanBuffer::new_unset(array.len())
            };
            return Some(BooleanArray::new(bool_array, array.nulls().cloned()));
        }
        Err(UnsupportedExpression::Expr) | Err(UnsupportedExpression::Op) => {
            return None;
        }
    };
    let op = expr.op();
    let needle = expr.literal();
    if let ByteViewOperator::SubString(_substring_op) = op
        && array.string_fingerprints.as_ref().is_none()
    {
        return None;
    }
    Some(array.compare_with(needle, op))
}

pub(super) async fn try_eval_predicate_on_disk(
    expr: &Arc<dyn PhysicalExpr>,
    array: &LiquidByteViewArray<DiskBuffer>,
) -> Option<BooleanArray> {
    let cur_expr = match ByteViewExpression::try_from(expr) {
        Ok(expr) => expr,
        Err(UnsupportedExpression::Constant(v)) => {
            let bool_array = if v {
                BooleanBuffer::new_set(array.len())
            } else {
                BooleanBuffer::new_unset(array.len())
            };
            return Some(BooleanArray::new(bool_array, array.nulls().cloned()));
        }
        Err(UnsupportedExpression::Expr) | Err(UnsupportedExpression::Op) => {
            return None;
        }
    };

    let op = cur_expr.op();
    let needle = cur_expr.literal();

    if let ByteViewOperator::SubString(_substring_op) = op
        && array.string_fingerprints.as_ref().is_none()
    {
        return None;
    }
    let result = array.compare_with(needle, op).await;
    Some(result)
}

use std::fmt::Display;

/// Detailed memory usage of the byte view array
pub struct ByteViewArrayMemoryUsage {
    /// Memory usage of the dictionary key
    pub dictionary_key: usize,
    /// Memory usage of the prefix keys
    pub prefix_keys: usize,
    /// Memory usage of the raw FSST buffer
    pub fsst_buffer: usize,
    /// Memory usage of the shared prefix
    pub shared_prefix: usize,
    /// Memory usage of the string fingerprints
    pub string_fingerprints: usize,
    /// Memory usage of the struct size
    pub struct_size: usize,
}

impl Display for ByteViewArrayMemoryUsage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ByteViewArrayMemoryUsage")
            .field("dictionary_key", &self.dictionary_key)
            .field("prefix_keys", &self.prefix_keys)
            .field("fsst_buffer", &self.fsst_buffer)
            .field("shared_prefix", &self.shared_prefix)
            .field("string_fingerprints", &self.string_fingerprints)
            .field("struct_size", &self.struct_size)
            .field("total", &self.total())
            .finish()
    }
}

impl ByteViewArrayMemoryUsage {
    /// Get the total memory usage of the byte view array
    pub fn total(&self) -> usize {
        self.dictionary_key
            + self.prefix_keys
            + self.fsst_buffer
            + self.shared_prefix
            + self.string_fingerprints
            + self.struct_size
    }
}

impl std::ops::AddAssign for ByteViewArrayMemoryUsage {
    fn add_assign(&mut self, other: Self) {
        self.dictionary_key += other.dictionary_key;
        self.prefix_keys += other.prefix_keys;
        self.fsst_buffer += other.fsst_buffer;
        self.shared_prefix += other.shared_prefix;
        self.string_fingerprints += other.string_fingerprints;
        self.struct_size += other.struct_size;
    }
}