vortex_array/arrays/dict/

array.rs

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

use std::fmt::Display;
use std::fmt::Formatter;

use num_traits::AsPrimitive;
use smallvec::smallvec;
use vortex_buffer::BitBuffer;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_mask::AllOr;

use crate::ArrayRef;
use crate::ArraySlots;
use crate::ExecutionCtx;
use crate::array::Array;
use crate::array::ArrayParts;
use crate::array::TypedArrayRef;
use crate::array_slots;
use crate::arrays::Dict;
use crate::arrays::PrimitiveArray;
use crate::dtype::DType;
use crate::dtype::PType;
use crate::match_each_integer_ptype;

#[derive(Clone, prost::Message)]
pub struct DictMetadata {
    #[prost(uint32, tag = "1")]
    pub(super) values_len: u32,
    #[prost(enumeration = "PType", tag = "2")]
    pub(super) codes_ptype: i32,
    // nullable codes are optional since they were added after stabilisation.
    #[prost(optional, bool, tag = "3")]
    pub(super) is_nullable_codes: Option<bool>,
    // all_values_referenced is optional for backward compatibility.
    // true = all dictionary values are definitely referenced by at least one code.
    // false/None = unknown whether all values are referenced (conservative default).
    #[prost(optional, bool, tag = "4")]
    pub(super) all_values_referenced: Option<bool>,
}

#[array_slots(Dict)]
pub struct DictSlots {
    /// The codes array mapping each element to a dictionary entry.
    #[slot(0)]
    pub codes: ArrayRef,
    /// The dictionary values array containing the unique values.
    #[slot(1)]
    pub values: ArrayRef,
}

#[derive(Debug, Clone)]
pub struct DictData {
    /// Indicates whether all dictionary values are definitely referenced by at least one code.
    /// `true` = all values are referenced (computed during encoding).
    /// `false` = unknown/might have unreferenced values.
    /// In case this is incorrect never use this to enable memory unsafe behaviour just semantically
    /// incorrect behaviour.
    pub(super) all_values_referenced: bool,
}

impl Display for DictData {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "all_values_referenced: {}", self.all_values_referenced)
    }
}

impl DictData {
    /// Build a new `DictArray` without validating the codes or values.
    ///
    /// # Safety
    /// This should be called only when you can guarantee the invariants checked
    /// by the safe `DictArray::try_new` constructor are valid, for example when
    /// you are filtering or slicing an existing valid `DictArray`.
    pub unsafe fn new_unchecked() -> Self {
        Self {
            all_values_referenced: false,
        }
    }

    /// Set whether all dictionary values are definitely referenced.
    ///
    /// # Safety
    /// The caller must ensure that when setting `all_values_referenced = true`, ALL dictionary
    /// values are actually referenced by at least one valid code. Setting this incorrectly can
    /// lead to incorrect query results in operations like min/max.
    ///
    /// This is typically only set to `true` during dictionary encoding when we know for certain
    /// that all values are referenced.
    pub unsafe fn set_all_values_referenced(mut self, all_values_referenced: bool) -> Self {
        self.all_values_referenced = all_values_referenced;
        self
    }

    /// Build a new `DictArray` from its components, `codes` and `values`.
    ///
    /// This constructor will panic if `codes` or `values` do not pass validation for building
    /// a new `DictArray`. See `DictArray::try_new` for a description of the error conditions.
    pub fn new(codes_dtype: &DType) -> Self {
        Self::try_new(codes_dtype).vortex_expect("DictArray new")
    }

    /// Build a new `DictArray` from its components, `codes` and `values`.
    ///
    /// The codes must be integers, and may be nullable. Values can be any
    /// type, and may also be nullable. This mirrors the nullability of the Arrow `DictionaryArray`.
    ///
    /// # Errors
    ///
    /// The `codes` **must** be integers, and the maximum code must be less than the length
    /// of the `values` array. Otherwise, this constructor returns an error.
    ///
    /// It is an error to provide a nullable `codes` with non-nullable `values`.
    pub(crate) fn try_new(codes_dtype: &DType) -> VortexResult<Self> {
        if !codes_dtype.is_int() {
            vortex_bail!(MismatchedTypes: "int", codes_dtype);
        }

        Ok(unsafe { Self::new_unchecked() })
    }
}

pub trait DictArrayExt: TypedArrayRef<Dict> + DictArraySlotsExt {
    #[inline]
    fn has_all_values_referenced(&self) -> bool {
        self.all_values_referenced
    }

    fn validate_all_values_referenced(&self, ctx: &mut ExecutionCtx) -> VortexResult<()> {
        if self.has_all_values_referenced() {
            if !self.codes().is_host() {
                return Ok(());
            }

            let referenced_mask = self.compute_referenced_values_mask(true, ctx)?;
            let all_referenced = referenced_mask.true_count() == referenced_mask.len();

            vortex_ensure!(all_referenced, "value in dict not referenced");
        }

        Ok(())
    }

    fn compute_referenced_values_mask(
        &self,
        referenced: bool,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<BitBuffer> {
        let codes = self.codes();
        let codes_validity = codes.validity()?.execute_mask(codes.len(), ctx)?;
        let codes_primitive = codes.clone().execute::<PrimitiveArray>(ctx)?;
        let values_len = self.values().len();

        let mut values_vec = vec![!referenced; values_len];
        let saturated = match codes_validity.bit_buffer() {
            AllOr::All => match_each_integer_ptype!(codes_primitive.ptype(), |P| {
                mark_codes(&mut values_vec, referenced, codes_primitive.as_slice::<P>())
            }),
            AllOr::None => false,
            AllOr::Some(mask) => match_each_integer_ptype!(codes_primitive.ptype(), |P| {
                mark_masked_codes(
                    &mut values_vec,
                    referenced,
                    codes_primitive.as_slice::<P>(),
                    mask,
                    codes_validity.true_count(),
                )
            }),
        };
        if saturated {
            return Ok(BitBuffer::full(referenced, values_len));
        }

        Ok(BitBuffer::from(values_vec))
    }
}
impl<T: TypedArrayRef<Dict>> DictArrayExt for T {}

/// Number of codes visited between checks for whether every dictionary value is referenced.
const EARLY_EXIT_CHUNK: usize = 2048;

/// Sets `values[code]` to `referenced` for every code, returning `true` early once every entry
/// of `values` is `referenced`, since further codes cannot change the result.
///
/// With fewer codes than values not every value can be referenced, so the saturation checks are
/// skipped entirely.
fn mark_codes<P: AsPrimitive<usize>>(values: &mut [bool], referenced: bool, codes: &[P]) -> bool {
    if codes.len() < values.len() {
        for code in codes {
            values[code.as_()] = referenced;
        }
        return false;
    }

    let mut cursor = 0;
    for chunk in codes.chunks(EARLY_EXIT_CHUNK) {
        for code in chunk {
            values[code.as_()] = referenced;
        }
        if is_saturated(values, &mut cursor, referenced) {
            return true;
        }
    }
    false
}

/// Like [`mark_codes`], but only for the codes at the `valid_count` set positions of `mask`.
fn mark_masked_codes<P: AsPrimitive<usize>>(
    values: &mut [bool],
    referenced: bool,
    codes: &[P],
    mask: &BitBuffer,
    valid_count: usize,
) -> bool {
    if valid_count < values.len() {
        mask.set_indices().for_each(|idx| {
            values[codes[idx].as_()] = referenced;
        });
        return false;
    }

    let mut cursor = 0;
    for start in (0..codes.len()).step_by(EARLY_EXIT_CHUNK) {
        let end = (start + EARLY_EXIT_CHUNK).min(codes.len());
        let codes = &codes[start..end];
        mask.slice(start..end).set_indices().for_each(|idx| {
            values[codes[idx].as_()] = referenced;
        });
        if is_saturated(values, &mut cursor, referenced) {
            return true;
        }
    }
    false
}

/// Advances `cursor` past the entries of `values` equal to `referenced` and returns whether that
/// is all of them. Callers share one cursor that only moves forward, so repeated checks cost
/// `O(values.len())` in total.
#[inline]
fn is_saturated(values: &[bool], cursor: &mut usize, referenced: bool) -> bool {
    // Skip whole blocks with a branch-free comparison the compiler can vectorize, then find the
    // exact stopping point within the first block that is not fully marked.
    const BLOCK: usize = 32;
    let rest = &values[*cursor..];
    let full_blocks = rest
        .as_chunks::<BLOCK>()
        .0
        .iter()
        .take_while(|block| block.iter().fold(true, |acc, &v| acc & (v == referenced)))
        .count();
    *cursor += full_blocks * BLOCK;
    *cursor += values[*cursor..]
        .iter()
        .take_while(|&&v| v == referenced)
        .count();
    *cursor == values.len()
}

/// Concrete parts of a [`DictArray`](super::DictArray) after iterative execution.
pub struct DictParts {
    pub dtype: DType,
    pub codes: ArrayRef,
    pub values: ArrayRef,
}

pub trait DictOwnedExt {
    fn into_parts(self) -> DictParts;
}

impl DictOwnedExt for Array<Dict> {
    fn into_parts(self) -> DictParts {
        match self.try_into_parts() {
            Ok(array_parts) => {
                let slots = DictSlots::from_slots(array_parts.slots);
                DictParts {
                    dtype: array_parts.dtype,
                    codes: slots.codes,
                    values: slots.values,
                }
            }
            Err(array) => {
                let slots = DictSlotsView::from_slots(array.slots());
                DictParts {
                    dtype: array.dtype().clone(),
                    codes: slots.codes.clone(),
                    values: slots.values.clone(),
                }
            }
        }
    }
}

impl Array<Dict> {
    /// Build a new `DictArray` from its components, `codes` and `values`.
    pub fn new(codes: ArrayRef, values: ArrayRef) -> Self {
        Self::try_new(codes, values).vortex_expect("DictArray new")
    }

    /// Build a new `DictArray` from its components, `codes` and `values`.
    pub fn try_new(codes: ArrayRef, values: ArrayRef) -> VortexResult<Self> {
        let dtype = values
            .dtype()
            .union_nullability(codes.dtype().nullability());
        let len = codes.len();
        let data = DictData::try_new(codes.dtype())?;
        Array::try_from_parts(ArrayParts::new(
            Dict,
            dtype,
            len,
            data,
            smallvec![Some(codes), Some(values)],
        ))
    }

    /// Build a new `DictArray` without validating the codes or values.
    ///
    /// # Safety
    ///
    /// See [`DictData::new_unchecked`].
    pub unsafe fn new_unchecked(codes: ArrayRef, values: ArrayRef) -> Self {
        let dtype = values
            .dtype()
            .union_nullability(codes.dtype().nullability());
        let len = codes.len();
        let data = unsafe { DictData::new_unchecked() };
        unsafe {
            Array::from_parts_unchecked(ArrayParts::new(
                Dict,
                dtype,
                len,
                data,
                smallvec![Some(codes), Some(values)],
            ))
        }
    }

    /// Set whether all values in the dictionary are referenced by at least one code.
    ///
    /// # Safety
    ///
    /// See [`DictData::set_all_values_referenced`].
    pub unsafe fn set_all_values_referenced(self, all_values_referenced: bool) -> Self {
        let dtype = self.dtype().clone();
        let len = self.len();
        let slots: ArraySlots = self.slots().iter().cloned().collect();
        let data = unsafe {
            self.into_data()
                .set_all_values_referenced(all_values_referenced)
        };
        unsafe { Array::from_parts_unchecked(ArrayParts::new(Dict, dtype, len, data, slots)) }
    }
}

#[cfg(test)]
mod test {
    use rand::RngExt;
    use rand::SeedableRng;
    use rand::distr::Distribution;
    use rand::distr::StandardUniform;
    use rand::prelude::StdRng;
    use vortex_buffer::BitBuffer;
    use vortex_buffer::buffer;
    use vortex_error::VortexExpect;
    use vortex_error::VortexResult;
    use vortex_error::vortex_panic;
    use vortex_mask::AllOr;

    use crate::ArrayRef;
    use crate::IntoArray;
    use crate::VortexSessionExecute;
    use crate::array_session;
    use crate::arrays::ChunkedArray;
    use crate::arrays::DictArray;
    use crate::arrays::PrimitiveArray;
    use crate::arrays::VarBinViewArray;
    use crate::arrays::dict::DictArrayExt;
    use crate::assert_arrays_eq;
    use crate::builders::VarBinBuilder;
    use crate::builders::builder_with_capacity_in;
    use crate::dtype::DType;
    use crate::dtype::NativePType;
    use crate::dtype::Nullability::NonNullable;
    use crate::dtype::PType;
    use crate::dtype::UnsignedPType;
    use crate::validity::Validity;

    #[test]
    fn nullable_codes_validity() {
        let dict = DictArray::try_new(
            PrimitiveArray::new(
                buffer![0u32, 1, 2, 2, 1],
                Validity::from(BitBuffer::from(vec![true, false, true, false, true])),
            )
            .into_array(),
            PrimitiveArray::new(buffer![3, 6, 9], Validity::AllValid).into_array(),
        )
        .unwrap();
        let mask = dict
            .as_ref()
            .validity()
            .unwrap()
            .execute_mask(
                dict.as_ref().len(),
                &mut array_session().create_execution_ctx(),
            )
            .unwrap();
        let AllOr::Some(indices) = mask.indices() else {
            vortex_panic!("Expected indices from mask")
        };
        assert_eq!(indices, [0, 2, 4]);
    }

    #[test]
    fn nullable_values_validity() {
        let dict = DictArray::try_new(
            buffer![0u32, 1, 2, 2, 1].into_array(),
            PrimitiveArray::new(
                buffer![3, 6, 9],
                Validity::from(BitBuffer::from(vec![true, false, false])),
            )
            .into_array(),
        )
        .unwrap();
        let mask = dict
            .as_ref()
            .validity()
            .unwrap()
            .execute_mask(
                dict.as_ref().len(),
                &mut array_session().create_execution_ctx(),
            )
            .unwrap();
        let AllOr::Some(indices) = mask.indices() else {
            vortex_panic!("Expected indices from mask")
        };
        assert_eq!(indices, [0]);
    }

    #[test]
    fn nullable_codes_and_values() {
        let dict = DictArray::try_new(
            PrimitiveArray::new(
                buffer![0u32, 1, 2, 2, 1],
                Validity::from(BitBuffer::from(vec![true, false, true, false, true])),
            )
            .into_array(),
            PrimitiveArray::new(
                buffer![3, 6, 9],
                Validity::from(BitBuffer::from(vec![false, true, true])),
            )
            .into_array(),
        )
        .unwrap();
        let mask = dict
            .as_ref()
            .validity()
            .unwrap()
            .execute_mask(
                dict.as_ref().len(),
                &mut array_session().create_execution_ctx(),
            )
            .unwrap();
        let AllOr::Some(indices) = mask.indices() else {
            vortex_panic!("Expected indices from mask")
        };
        assert_eq!(indices, [2, 4]);
    }

    #[test]
    fn nullable_codes_and_non_null_values() {
        let dict = DictArray::try_new(
            PrimitiveArray::new(
                buffer![0u32, 1, 2, 2, 1],
                Validity::from(BitBuffer::from(vec![true, false, true, false, true])),
            )
            .into_array(),
            PrimitiveArray::new(buffer![3, 6, 9], Validity::NonNullable).into_array(),
        )
        .unwrap();
        let mask = dict
            .as_ref()
            .validity()
            .unwrap()
            .execute_mask(
                dict.as_ref().len(),
                &mut array_session().create_execution_ctx(),
            )
            .unwrap();
        let AllOr::Some(indices) = mask.indices() else {
            vortex_panic!("Expected indices from mask")
        };
        assert_eq!(indices, [0, 2, 4]);
    }

    fn make_dict_primitive_chunks<T: NativePType, Code: UnsignedPType>(
        len: usize,
        unique_values: usize,
        chunk_count: usize,
    ) -> (ArrayRef, Vec<T>)
    where
        StandardUniform: Distribution<T>,
    {
        let mut rng = StdRng::seed_from_u64(0);

        let mut expected = Vec::with_capacity(len * chunk_count);
        let array = (0..chunk_count)
            .map(|_| {
                let values = (0..unique_values)
                    .map(|_| rng.random::<T>())
                    .collect::<PrimitiveArray>();
                let codes = (0..len)
                    .map(|_| {
                        Code::from(rng.random_range(0..unique_values)).vortex_expect("valid value")
                    })
                    .collect::<PrimitiveArray>();

                expected.extend(codes.as_slice::<Code>().iter().map(|code| {
                    values.as_slice::<T>()[<Code as num_traits::AsPrimitive<usize>>::as_(*code)]
                }));
                DictArray::try_new(codes.into_array(), values.into_array())
                    .vortex_expect("DictArray creation should succeed in arbitrary impl")
                    .into_array()
            })
            .collect::<ChunkedArray>()
            .into_array();
        (array, expected)
    }

    #[test]
    fn test_dict_utf8_append_to_varbin_builder() -> VortexResult<()> {
        let values = VarBinViewArray::from_iter_str(["zero", "one", "two"]);
        let dict = DictArray::try_new(buffer![2u8, 0, 2, 1].into_array(), values.into_array())?;
        let expected = VarBinViewArray::from_iter_str(["two", "zero", "two", "one"]);
        let mut ctx = array_session().create_execution_ctx();
        let mut builder = VarBinBuilder::<i32>::with_capacity_in(
            dict.dtype().clone(),
            dict.len(),
            ctx.allocator(),
        );

        dict.into_array()
            .append_to_builder(&mut builder, &mut ctx)?;

        assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx);
        Ok(())
    }

    #[test]
    fn test_dict_array_from_primitive_chunks() -> VortexResult<()> {
        let mut ctx = array_session().create_execution_ctx();
        let len = 2;
        let chunk_count = 2;
        let (array, expected) = make_dict_primitive_chunks::<u64, u64>(len, 2, chunk_count);

        let mut builder = builder_with_capacity_in(
            &DType::Primitive(PType::U64, NonNullable),
            len * chunk_count,
            ctx.allocator(),
        );
        array.append_to_builder(
            builder.as_mut(),
            &mut array_session().create_execution_ctx(),
        )?;

        let prim_into = builder.finish_into_canonical(&mut ctx).into_primitive();

        let expected = PrimitiveArray::from_iter(expected);
        assert_arrays_eq!(array, expected, &mut ctx);
        assert_arrays_eq!(prim_into, expected, &mut ctx);
        Ok(())
    }

    #[cfg_attr(miri, ignore)]
    #[test]
    fn test_dict_metadata() {
        use prost::Message;

        use super::DictMetadata;
        use crate::test_harness::check_metadata;

        check_metadata(
            "dict.metadata",
            &DictMetadata {
                codes_ptype: PType::U64 as i32,
                values_len: u32::MAX,
                is_nullable_codes: None,
                all_values_referenced: None,
            }
            .encode_to_vec(),
        );
    }

    /// Covers the early exit taken once every value is referenced, including when that happens
    /// after the first chunk, alongside cases that never saturate.
    #[rstest::rstest]
    #[case::saturates_in_first_chunk(16, 5000, 16, None)]
    #[case::saturates_in_later_chunk(4096, 60_000, 4096, None)]
    #[case::saturates_unaligned_to_scan_block(4001, 60_000, 4001, None)]
    #[case::never_saturates(1024, 10_000, 1000, None)]
    #[case::fewer_codes_than_values(1024, 1000, 1024, None)]
    #[case::fewer_valid_codes_than_values(1024, 1500, 1024, Some(0.5))]
    #[case::saturates_with_nulls(64, 10_000, 64, Some(0.5))]
    #[case::never_saturates_with_nulls(64, 10_000, 63, Some(0.5))]
    fn referenced_values_mask_matches_reference(
        #[case] num_values: u32,
        #[case] num_codes: usize,
        #[case] referenced_range: u32,
        #[case] valid_fraction: Option<f64>,
        #[values(true, false)] referenced: bool,
    ) -> VortexResult<()> {
        let mut rng = StdRng::seed_from_u64(0);
        let codes: Vec<Option<u32>> = (0..num_codes)
            .map(|_| {
                valid_fraction
                    .is_none_or(|f| rng.random_bool(f))
                    .then(|| rng.random_range(0..referenced_range))
            })
            .collect();

        let mut expected = vec![!referenced; num_values as usize];
        for code in codes.iter().flatten() {
            expected[*code as usize] = referenced;
        }

        let codes = match valid_fraction {
            Some(_) => PrimitiveArray::from_option_iter(codes).into_array(),
            None => PrimitiveArray::from_iter(codes.into_iter().flatten()).into_array(),
        };
        let dict =
            DictArray::try_new(codes, PrimitiveArray::from_iter(0..num_values).into_array())?;
        let mut ctx = array_session().create_execution_ctx();
        let mask = dict.compute_referenced_values_mask(referenced, &mut ctx)?;

        assert_eq!(mask, BitBuffer::from(expected));
        Ok(())
    }
}