vortex-runend-bool 0.21.1

mod invert;

use arrow_buffer::BooleanBuffer;
use vortex_array::array::BoolArray;
use vortex_array::compute::{slice, ComputeVTable, InvertFn, ScalarAtFn, SliceFn, TakeFn};
use vortex_array::variants::PrimitiveArrayTrait;
use vortex_array::{ArrayDType, ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant};
use vortex_dtype::match_each_integer_ptype;
use vortex_error::{vortex_bail, VortexResult};
use vortex_scalar::Scalar;

use crate::{value_at_index, RunEndBoolArray, RunEndBoolEncoding};

impl ComputeVTable for RunEndBoolEncoding {
    fn invert_fn(&self) -> Option<&dyn InvertFn<ArrayData>> {
        Some(self)
    }

    fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
        Some(self)
    }
    fn slice_fn(&self) -> Option<&dyn SliceFn<ArrayData>> {
        Some(self)
    }
    fn take_fn(&self) -> Option<&dyn TakeFn<ArrayData>> {
        Some(self)
    }
}

impl ScalarAtFn<RunEndBoolArray> for RunEndBoolEncoding {
    fn scalar_at(&self, array: &RunEndBoolArray, index: usize) -> VortexResult<Scalar> {
        let start = array.start();
        Ok(Scalar::bool(
            value_at_index(array.find_physical_index(index)?, start),
            array.dtype().nullability(),
        ))
    }
}

impl TakeFn<RunEndBoolArray> for RunEndBoolEncoding {
    fn take(&self, array: &RunEndBoolArray, indices: &ArrayData) -> VortexResult<ArrayData> {
        let primitive_indices = indices.clone().into_primitive()?;
        let physical_indices = match_each_integer_ptype!(primitive_indices.ptype(), |$P| {
            primitive_indices
                .into_maybe_null_slice::<$P>()
                .into_iter()
                .map(|idx| idx as usize)
                .map(|idx| {
                    if idx >= array.len() {
                        vortex_bail!(OutOfBounds: idx, 0, array.len())
                    }
                    array.find_physical_index(idx)
                })
                .collect::<VortexResult<Vec<_>>>()?
        });
        let start = array.start();
        BoolArray::try_new(
            BooleanBuffer::from_iter(
                physical_indices
                    .into_iter()
                    .map(|it| value_at_index(it, start)),
            ),
            array.validity().take(indices)?,
        )
        .map(|a| a.into_array())
    }
}

impl SliceFn<RunEndBoolArray> for RunEndBoolEncoding {
    fn slice(&self, array: &RunEndBoolArray, start: usize, stop: usize) -> VortexResult<ArrayData> {
        let new_length = stop - start;

        let (slice_begin, slice_end) = if new_length == 0 {
            let ends_len = array.ends().len();
            (ends_len, ends_len)
        } else {
            let physical_begin = array.find_physical_index(start)?;
            let physical_end = array.find_physical_index(stop)?;
            (physical_begin, physical_end + 1)
        };

        Ok(RunEndBoolArray::with_offset_and_size(
            slice(array.ends(), slice_begin, slice_end)?,
            value_at_index(slice_begin, array.start()),
            array.validity().slice(start, stop)?,
            new_length,
            if new_length == 0 {
                0
            } else {
                start + array.offset()
            },
        )?
        .into_array())
    }
}

#[cfg(test)]
mod tests {
    use arrow_buffer::BooleanBuffer;
    use vortex_array::array::PrimitiveArray;
    use vortex_array::compute::{scalar_at, slice, take};
    use vortex_array::validity::Validity;
    use vortex_array::{ArrayDType, ArrayLen, IntoArrayData, IntoArrayVariant};
    use vortex_dtype::Nullability;
    use vortex_scalar::Scalar;

    use crate::RunEndBoolArray;

    #[test]
    fn slice_at_end() {
        let re_array =
            RunEndBoolArray::try_new(vec![7_u64, 10].into_array(), false, Validity::NonNullable)
                .unwrap();

        assert_eq!(re_array.len(), 10);

        let sliced_array = slice(&re_array, re_array.len(), re_array.len()).unwrap();
        assert!(sliced_array.is_empty());

        let re_slice = RunEndBoolArray::try_from(sliced_array).unwrap();
        assert!(re_slice.ends().is_empty());
    }

    #[test]
    fn scalar_at_nullability() {
        let re_array =
            RunEndBoolArray::try_new(vec![7_u64, 10].into_array(), false, Validity::AllValid)
                .unwrap();

        assert_eq!(
            scalar_at(&re_array, 0).unwrap(),
            Scalar::bool(false, Nullability::Nullable)
        );
    }

    #[test]
    fn take_nullable() {
        let re_array = RunEndBoolArray::try_new(
            vec![7_u64, 10].into_array(),
            false,
            Validity::from(BooleanBuffer::from(vec![
                false, false, true, true, true, true, true, true, false, false,
            ])),
        )
        .unwrap();

        let taken = take(&re_array, PrimitiveArray::from(vec![6, 9])).unwrap();
        let taken_bool = taken.into_bool().unwrap();
        assert_eq!(taken_bool.dtype(), re_array.dtype());
        assert_eq!(
            taken_bool.boolean_buffer(),
            BooleanBuffer::from(vec![false, true])
        );
    }
}