vortex-array 0.88.0

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

//! Encoding-specific mask reduction and execution adapters.

use vortex_error::VortexResult;
use vortex_error::vortex_err;

use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::array::ArrayView;
use crate::array::VTable;
use crate::arrays::Bool;
use crate::arrays::Constant;
use crate::arrays::scalar_fn::ExactScalarFn;
use crate::arrays::scalar_fn::ScalarFnArrayView;
use crate::kernel::ExecuteParentKernel;
use crate::optimizer::rules::ArrayParentReduceRule;
use crate::scalar_fn::fns::mask::Mask as MaskExpr;

/// Mask an array without reading buffers.
///
/// This trait is for mask implementations that can operate purely on array metadata and
/// structure without needing to read or execute on the underlying buffers. Implementations
/// should return `None` if masking requires buffer access.
///
/// The `mask` parameter is a boolean array where true=keep/valid, false=null-out.
///
/// # Preconditions
///
/// The mask has non-nullable Boolean dtype and the same length as the array. It may be lazy.
pub trait MaskReduce: VTable {
    /// Whether reading this encoding's validity requires only metadata access.
    ///
    /// Enables lazy masks on all-valid inputs. Only opt in when `array.validity()` cannot read
    /// buffers or execute children, including for nullable inputs.
    const VALIDITY_IS_METADATA_ONLY: bool = false;

    fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
}

/// Mask an array, potentially reading buffers.
///
/// Unlike [`MaskReduce`], this trait is for mask implementations that may need to read
/// and execute on the underlying buffers to produce the masked result.
///
/// The `mask` parameter is a boolean array where true=keep/valid, false=null-out.
///
/// # Preconditions
///
/// The mask is guaranteed to have the same length as the array. Trivial cases
/// (`AllValid`, `AllInvalid`, `NonNullable`) are handled by the caller before dispatch.
pub trait MaskKernel: VTable {
    fn mask(
        array: ArrayView<'_, Self>,
        mask: &ArrayRef,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<ArrayRef>>;
}

/// Adaptor that wraps a [`MaskReduce`] impl as an [`ArrayParentReduceRule`].
#[derive(Default, Debug)]
pub struct MaskReduceAdaptor<V>(pub V);

impl<V> ArrayParentReduceRule<V> for MaskReduceAdaptor<V>
where
    V: MaskReduce,
{
    type Parent = ExactScalarFn<MaskExpr>;

    fn reduce_parent(
        &self,
        array: ArrayView<'_, V>,
        parent: ScalarFnArrayView<'_, MaskExpr>,
        child_idx: usize,
    ) -> VortexResult<Option<ArrayRef>> {
        // Only reduce the input child (index 0), not the mask child (index 1).
        if child_idx != 0 {
            return Ok(None);
        }

        let parent_ref: ArrayRef = (*parent).clone();
        let mask_child = parent_ref
            .nth_child(1)
            .ok_or_else(|| vortex_err!("Mask expression must have 2 children"))?;

        if mask_child.as_opt::<Bool>().is_none() && mask_child.as_opt::<Constant>().is_none() {
            let can_attach_mask =
                V::VALIDITY_IS_METADATA_ONLY && array.validity()?.definitely_no_nulls();

            if !can_attach_mask {
                return Ok(None);
            }
        }

        <V as MaskReduce>::mask(array, mask_child)
    }
}

/// Adaptor that wraps a [`MaskKernel`] impl as an [`ExecuteParentKernel`].
#[derive(Default, Debug)]
pub struct MaskExecuteAdaptor<V>(pub V);

impl<V> ExecuteParentKernel<V> for MaskExecuteAdaptor<V>
where
    V: MaskKernel,
{
    type Parent = ExactScalarFn<MaskExpr>;

    fn execute_parent(
        &self,
        array: ArrayView<'_, V>,
        parent: ScalarFnArrayView<'_, MaskExpr>,
        child_idx: usize,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<ArrayRef>> {
        // Only execute the input child (index 0), not the mask child (index 1).
        if child_idx != 0 {
            return Ok(None);
        }
        let mask_child = parent
            .nth_child(1)
            .ok_or_else(|| vortex_err!("Mask expression must have 2 children"))?;
        <V as MaskKernel>::mask(array, mask_child, ctx)
    }
}

#[cfg(test)]
mod tests {
    use rstest::rstest;
    use vortex_buffer::buffer;
    use vortex_error::VortexResult;

    use crate::ArrayEq;
    use crate::ArrayRef;
    use crate::EqMode;
    use crate::IntoArray;
    use crate::arrays::BoolArray;
    use crate::arrays::ConstantArray;
    use crate::arrays::Primitive;
    use crate::arrays::PrimitiveArray;
    use crate::arrays::ScalarFn;
    use crate::arrays::scalar_fn::ExactScalarFn;
    use crate::assert_arrays_eq;
    use crate::builtins::ArrayBuiltins;
    use crate::dtype::Nullability;
    use crate::executor::VortexSessionExecute;
    use crate::optimizer::ArrayOptimizer;
    use crate::scalar::Scalar;
    use crate::scalar_fn::fns::binary::Binary;
    use crate::scalar_fn::fns::mask::Mask as MaskExpr;
    use crate::scalar_fn::fns::operators::Operator;
    use crate::validity::Validity;

    fn lazy_mask() -> VortexResult<ArrayRef> {
        Binary::try_new(
            BoolArray::from_iter([true, true, false]).into_array(),
            BoolArray::from_iter([true, false, true]).into_array(),
            Operator::And,
        )
        .map(IntoArray::into_array)
    }

    #[rstest]
    #[case::non_nullable(Validity::NonNullable)]
    #[case::all_valid(Validity::AllValid)]
    fn all_valid_input_attaches_lazy_mask(#[case] validity: Validity) -> VortexResult<()> {
        let input = PrimitiveArray::new(buffer![1i32, 2, 3], validity);
        let mask = lazy_mask()?;
        let output = input.clone().into_array().mask(mask.clone())?;
        let primitive = output.as_::<Primitive>();

        assert!(
            primitive
                .buffer_handle()
                .array_eq(input.buffer_handle(), EqMode::Ptr)
        );
        assert_eq!(output.dtype(), &input.dtype().as_nullable());
        assert_eq!(output.len(), input.len());
        let Validity::Array(validity) = output.validity()? else {
            panic!("mask must remain array-backed");
        };
        assert!(ArrayRef::ptr_eq(&validity, &mask));

        Ok(())
    }

    #[rstest]
    #[case::partially_valid(false)]
    #[case::all_true_bitmap(true)]
    fn array_backed_validity_keeps_lazy_mask(#[case] all_true: bool) -> VortexResult<()> {
        let validity = BoolArray::from_iter([true, all_true, true]).into_array();
        let input =
            PrimitiveArray::new(buffer![1i32, 2, 3], Validity::Array(validity)).into_array();
        let output = input.mask(lazy_mask()?)?;
        assert!(output.is::<ExactScalarFn<MaskExpr>>());

        Ok(())
    }

    /// A constant Boolean mask child must take the metadata-only reduction path (pushing into the
    /// input encoding) rather than surviving as a `ScalarFn` wrapper that falls through to
    /// execution. Asserting the optimized encoding makes this fail before the adaptor accepts
    /// `Constant` masks, not just verifying values that could pass through the execution fallback.
    #[rstest]
    #[case(true)]
    #[case(false)]
    fn constant_mask_reduces_into_input(#[case] mask_value: bool) -> VortexResult<()> {
        let input = buffer![1i32, 2, 3, 4, 5].into_array();
        let mask = ConstantArray::new(
            Scalar::bool(mask_value, Nullability::NonNullable),
            input.len(),
        )
        .into_array();

        let masked = MaskExpr::try_new(input, mask)?.into_array();
        assert!(
            masked.is::<ScalarFn>(),
            "expected an un-optimized ScalarFn wrapper before optimization"
        );

        let optimized = masked.optimize()?;
        assert!(
            !optimized.is::<ScalarFn>(),
            "constant mask should not fall through to execution, got {}",
            optimized.encoding_id()
        );
        assert!(
            optimized.is::<Primitive>(),
            "constant mask should reduce into the Primitive input, got {}",
            optimized.encoding_id()
        );

        let mut ctx = crate::array_session().create_execution_ctx();
        let expected = if mask_value {
            PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3), Some(4), Some(5)])
        } else {
            PrimitiveArray::from_option_iter([None::<i32>, None, None, None, None])
        };
        assert_arrays_eq!(optimized, expected, &mut ctx);

        Ok(())
    }
}