vortex-array 0.21.1

use vortex_dtype::field::Field;
use vortex_error::{vortex_err, VortexExpect, VortexResult};
use vortex_scalar::StructScalar;

use crate::array::sparse::SparseArray;
use crate::array::SparseEncoding;
use crate::variants::{
    BinaryArrayTrait, BoolArrayTrait, ExtensionArrayTrait, ListArrayTrait, NullArrayTrait,
    PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait, VariantsVTable,
};
use crate::{ArrayData, ArrayLen, IntoArrayData};

/// Sparse arrays support all DTypes
impl VariantsVTable<SparseArray> for SparseEncoding {
    fn as_null_array<'a>(&self, array: &'a SparseArray) -> Option<&'a dyn NullArrayTrait> {
        Some(array)
    }

    fn as_bool_array<'a>(&self, array: &'a SparseArray) -> Option<&'a dyn BoolArrayTrait> {
        Some(array)
    }

    fn as_primitive_array<'a>(
        &self,
        array: &'a SparseArray,
    ) -> Option<&'a dyn PrimitiveArrayTrait> {
        Some(array)
    }

    fn as_utf8_array<'a>(&self, array: &'a SparseArray) -> Option<&'a dyn Utf8ArrayTrait> {
        Some(array)
    }

    fn as_binary_array<'a>(&self, array: &'a SparseArray) -> Option<&'a dyn BinaryArrayTrait> {
        Some(array)
    }

    fn as_struct_array<'a>(&self, array: &'a SparseArray) -> Option<&'a dyn StructArrayTrait> {
        Some(array)
    }

    fn as_list_array<'a>(&self, array: &'a SparseArray) -> Option<&'a dyn ListArrayTrait> {
        Some(array)
    }

    fn as_extension_array<'a>(
        &self,
        array: &'a SparseArray,
    ) -> Option<&'a dyn ExtensionArrayTrait> {
        Some(array)
    }
}

impl NullArrayTrait for SparseArray {}

impl BoolArrayTrait for SparseArray {}

impl PrimitiveArrayTrait for SparseArray {}

impl Utf8ArrayTrait for SparseArray {}

impl BinaryArrayTrait for SparseArray {}

impl StructArrayTrait for SparseArray {
    fn field(&self, idx: usize) -> Option<ArrayData> {
        let new_patches = self
            .patches()
            .map_values_opt(|values| values.as_struct_array().and_then(|s| s.field(idx)))
            .vortex_expect("field array length should equal struct array length")?;
        let scalar = StructScalar::try_from(&self.fill_scalar())
            .ok()?
            .field_by_idx(idx)?;

        Some(
            SparseArray::try_new_from_patches(
                new_patches,
                self.len(),
                self.indices_offset(),
                scalar,
            )
            .ok()?
            .into_array(),
        )
    }

    fn project(&self, projection: &[Field]) -> VortexResult<ArrayData> {
        let new_patches = self.patches().map_values(|values| {
            values
                .as_struct_array()
                .ok_or_else(|| vortex_err!("Chunk was not a StructArray"))?
                .project(projection)
        })?;
        let scalar = StructScalar::try_from(&self.fill_scalar())?.project(projection)?;

        SparseArray::try_new_from_patches(new_patches, self.len(), self.indices_offset(), scalar)
            .map(IntoArrayData::into_array)
    }
}

impl ListArrayTrait for SparseArray {}

impl ExtensionArrayTrait for SparseArray {
    fn storage_data(&self) -> ArrayData {
        SparseArray::try_new_from_patches(
            self.patches()
                .map_values(|values| {
                    Ok(values
                        .as_extension_array()
                        .vortex_expect("Expected extension array")
                        .storage_data())
                })
                .vortex_expect("as_extension_array preserves the length"),
            self.len(),
            self.indices_offset(),
            self.fill_scalar(),
        )
        .vortex_expect("Failed to create new sparse array")
        .into_array()
    }
}

#[cfg(test)]
mod tests {
    use vortex_scalar::Scalar;

    use crate::array::{BoolArray, PrimitiveArray, SparseArray};
    use crate::compute::invert;
    use crate::{IntoArrayData, IntoArrayVariant};

    #[test]
    fn invert_bools_non_null_fill() {
        let sparse_bools = SparseArray::try_new(
            PrimitiveArray::from(vec![0u64]).into_array(),
            BoolArray::from_iter([false]).into_array(),
            2,
            Scalar::from(true),
        )
        .unwrap()
        .into_array();
        let inverted = invert(&sparse_bools).unwrap();
        assert_eq!(
            inverted
                .into_bool()
                .unwrap()
                .boolean_buffer()
                .iter()
                .collect::<Vec<_>>(),
            vec![true, false]
        );
    }
}