vortex_array/arrays/struct_/compute/

mod.rs

mod cast;
mod filter;
mod mask;

use itertools::Itertools;
use vortex_dtype::Nullability::NonNullable;
use vortex_error::VortexResult;
use vortex_scalar::Scalar;

use crate::arrays::StructVTable;
use crate::arrays::struct_::StructArray;
use crate::compute::{
    IsConstantKernel, IsConstantKernelAdapter, IsConstantOpts, MinMaxKernel, MinMaxKernelAdapter,
    MinMaxResult, TakeKernel, TakeKernelAdapter, fill_null, is_constant_opts, take,
};
use crate::validity::Validity;
use crate::vtable::ValidityHelper;
use crate::{Array, ArrayRef, IntoArray, register_kernel};

impl TakeKernel for StructVTable {
    fn take(&self, array: &StructArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
        // If the struct array is empty then the indices must be all null, otherwise it will access
        // an out of bounds element
        if array.is_empty() {
            return StructArray::try_new_with_dtype(
                array.fields().to_vec(),
                array.struct_fields().clone(),
                indices.len(),
                Validity::AllInvalid,
            )
            .map(StructArray::into_array);
        }
        // The validity is applied to the struct validity,
        let inner_indices = &fill_null(
            indices,
            &Scalar::default_value(indices.dtype().with_nullability(NonNullable)),
        )?;
        StructArray::try_new_with_dtype(
            array
                .fields()
                .iter()
                .map(|field| take(field, inner_indices))
                .try_collect()?,
            array.struct_fields().clone(),
            indices.len(),
            array.validity().take(indices)?,
        )
        .map(|a| a.into_array())
    }
}

register_kernel!(TakeKernelAdapter(StructVTable).lift());

impl MinMaxKernel for StructVTable {
    fn min_max(&self, _array: &StructArray) -> VortexResult<Option<MinMaxResult>> {
        // TODO(joe): Implement struct min max
        Ok(None)
    }
}

register_kernel!(MinMaxKernelAdapter(StructVTable).lift());

impl IsConstantKernel for StructVTable {
    fn is_constant(
        &self,
        array: &StructArray,
        opts: &IsConstantOpts,
    ) -> VortexResult<Option<bool>> {
        let children = array.children();
        if children.is_empty() {
            return Ok(None);
        }

        for child in children.iter() {
            match is_constant_opts(child, opts)? {
                // Un-determined
                None => return Ok(None),
                Some(false) => return Ok(Some(false)),
                Some(true) => {}
            }
        }

        Ok(Some(true))
    }
}

register_kernel!(IsConstantKernelAdapter(StructVTable).lift());

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use Nullability::{NonNullable, Nullable};
    use vortex_buffer::buffer;
    use vortex_dtype::{DType, FieldNames, Nullability, PType, StructFields};
    use vortex_mask::Mask;
    use vortex_scalar::Scalar;

    use crate::arrays::{BoolArray, BooleanBuffer, PrimitiveArray, StructArray, VarBinArray};
    use crate::compute::conformance::mask::test_mask;
    use crate::compute::{cast, filter, take};
    use crate::validity::Validity;
    use crate::{Array, IntoArray as _};

    #[test]
    fn filter_empty_struct() {
        let struct_arr =
            StructArray::try_new(vec![].into(), vec![], 10, Validity::NonNullable).unwrap();
        let mask = vec![
            false, true, false, true, false, true, false, true, false, true,
        ];
        let filtered = filter(struct_arr.as_ref(), &Mask::from_iter(mask)).unwrap();
        assert_eq!(filtered.len(), 5);
    }

    #[test]
    fn take_empty_struct() {
        let struct_arr =
            StructArray::try_new(vec![].into(), vec![], 10, Validity::NonNullable).unwrap();
        let indices = PrimitiveArray::from_option_iter([Some(1), None]);
        let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap();
        assert_eq!(taken.len(), 2);

        assert_eq!(
            taken.scalar_at(0).unwrap(),
            Scalar::struct_(
                DType::Struct(Arc::new(StructFields::new([].into(), vec![])), Nullable),
                vec![]
            )
        );
        assert_eq!(
            taken.scalar_at(1).unwrap(),
            Scalar::null(DType::Struct(
                Arc::new(StructFields::new([].into(), vec![])),
                Nullable
            ))
        );
    }

    #[test]
    fn take_field_struct() {
        let struct_arr =
            StructArray::from_fields(&[("a", PrimitiveArray::from_iter(0..10).to_array())])
                .unwrap();
        let indices = PrimitiveArray::from_option_iter([Some(1), None]);
        let taken = take(struct_arr.as_ref(), indices.as_ref()).unwrap();
        assert_eq!(taken.len(), 2);

        assert_eq!(
            taken.scalar_at(0).unwrap(),
            Scalar::struct_(
                struct_arr.dtype().union_nullability(Nullable),
                vec![Scalar::primitive(1, NonNullable)],
            )
        );
        assert_eq!(
            taken.scalar_at(1).unwrap(),
            Scalar::null(struct_arr.dtype().union_nullability(Nullable),)
        );
    }

    #[test]
    fn filter_empty_struct_with_empty_filter() {
        let struct_arr =
            StructArray::try_new(vec![].into(), vec![], 0, Validity::NonNullable).unwrap();
        let filtered = filter(struct_arr.as_ref(), &Mask::from_iter::<[bool; 0]>([])).unwrap();
        assert_eq!(filtered.len(), 0);
    }

    #[test]
    fn test_mask_empty_struct() {
        test_mask(
            StructArray::try_new(vec![].into(), vec![], 5, Validity::NonNullable)
                .unwrap()
                .as_ref(),
        );
    }

    #[test]
    fn test_mask_complex_struct() {
        let xs = buffer![0i64, 1, 2, 3, 4].into_array();
        let ys = VarBinArray::from_iter(
            [Some("a"), Some("b"), None, Some("d"), None],
            DType::Utf8(Nullable),
        )
        .into_array();
        let zs =
            BoolArray::from_iter([Some(true), Some(true), None, None, Some(false)]).into_array();

        test_mask(
            StructArray::try_new(
                ["xs".into(), "ys".into(), "zs".into()].into(),
                vec![
                    StructArray::try_new(
                        ["left".into(), "right".into()].into(),
                        vec![xs.clone(), xs],
                        5,
                        Validity::NonNullable,
                    )
                    .unwrap()
                    .into_array(),
                    ys,
                    zs,
                ],
                5,
                Validity::NonNullable,
            )
            .unwrap()
            .as_ref(),
        );
    }

    #[test]
    fn test_cast_empty_struct() {
        let array = StructArray::try_new(vec![].into(), vec![], 5, Validity::NonNullable)
            .unwrap()
            .into_array();
        let non_nullable_dtype =
            DType::Struct(Arc::from(StructFields::new([].into(), vec![])), NonNullable);
        let casted = cast(&array, &non_nullable_dtype).unwrap();
        assert_eq!(casted.dtype(), &non_nullable_dtype);

        let nullable_dtype =
            DType::Struct(Arc::from(StructFields::new([].into(), vec![])), Nullable);
        let casted = cast(&array, &nullable_dtype).unwrap();
        assert_eq!(casted.dtype(), &nullable_dtype);
    }

    #[test]
    fn test_cast_cannot_change_name_order() {
        let array = StructArray::try_new(
            ["xs".into(), "ys".into(), "zs".into()].into(),
            vec![
                buffer![1u8].into_array(),
                buffer![1u8].into_array(),
                buffer![1u8].into_array(),
            ],
            1,
            Validity::NonNullable,
        )
        .unwrap();

        let tu8 = DType::Primitive(PType::U8, NonNullable);

        let result = cast(
            array.as_ref(),
            &DType::Struct(
                Arc::from(StructFields::new(
                    FieldNames::from(["ys".into(), "xs".into(), "zs".into()]),
                    vec![tu8.clone(), tu8.clone(), tu8],
                )),
                NonNullable,
            ),
        );
        assert!(
            result.as_ref().is_err_and(|err| {
                err.to_string()
                    .contains("cannot cast {xs=u8, ys=u8, zs=u8} to {ys=u8, xs=u8, zs=u8}")
            }),
            "{result:?}"
        );
    }

    #[test]
    fn test_cast_complex_struct() {
        let xs = PrimitiveArray::from_option_iter([Some(0i64), Some(1), Some(2), Some(3), Some(4)]);
        let ys = VarBinArray::from_vec(vec!["a", "b", "c", "d", "e"], DType::Utf8(Nullable));
        let zs = BoolArray::new(
            BooleanBuffer::from_iter([true, true, false, false, true]),
            Validity::AllValid,
        );
        let fully_nullable_array = StructArray::try_new(
            ["xs".into(), "ys".into(), "zs".into()].into(),
            vec![
                StructArray::try_new(
                    ["left".into(), "right".into()].into(),
                    vec![xs.to_array(), xs.to_array()],
                    5,
                    Validity::AllValid,
                )
                .unwrap()
                .into_array(),
                ys.into_array(),
                zs.into_array(),
            ],
            5,
            Validity::AllValid,
        )
        .unwrap()
        .into_array();

        let top_level_non_nullable = fully_nullable_array.dtype().as_nonnullable();
        let casted = cast(&fully_nullable_array, &top_level_non_nullable).unwrap();
        assert_eq!(casted.dtype(), &top_level_non_nullable);

        let non_null_xs_right = DType::Struct(
            Arc::from(StructFields::new(
                ["xs".into(), "ys".into(), "zs".into()].into(),
                vec![
                    DType::Struct(
                        Arc::from(StructFields::new(
                            ["left".into(), "right".into()].into(),
                            vec![
                                DType::Primitive(PType::I64, NonNullable),
                                DType::Primitive(PType::I64, Nullable),
                            ],
                        )),
                        Nullable,
                    ),
                    DType::Utf8(Nullable),
                    DType::Bool(Nullable),
                ],
            )),
            Nullable,
        );
        let casted = cast(&fully_nullable_array, &non_null_xs_right).unwrap();
        assert_eq!(casted.dtype(), &non_null_xs_right);

        let non_null_xs = DType::Struct(
            Arc::from(StructFields::new(
                ["xs".into(), "ys".into(), "zs".into()].into(),
                vec![
                    DType::Struct(
                        Arc::from(StructFields::new(
                            ["left".into(), "right".into()].into(),
                            vec![
                                DType::Primitive(PType::I64, Nullable),
                                DType::Primitive(PType::I64, Nullable),
                            ],
                        )),
                        NonNullable,
                    ),
                    DType::Utf8(Nullable),
                    DType::Bool(Nullable),
                ],
            )),
            Nullable,
        );
        let casted = cast(&fully_nullable_array, &non_null_xs).unwrap();
        assert_eq!(casted.dtype(), &non_null_xs);
    }
}