vortex-array 0.59.4

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

//! Compute kernels on top of Vortex Arrays.
//!
//! We aim to provide a basic set of compute kernels that can be used to efficiently index, slice,
//! and filter Vortex Arrays in their encoded forms.
//!
//! Every array encoding has the ability to implement their own efficient implementations of these
//! operators, else we will decode, and perform the equivalent operator from Arrow.

use std::any::Any;
use std::any::type_name;
use std::fmt::Debug;
use std::fmt::Formatter;

use arcref::ArcRef;
pub use boolean::*;
#[expect(deprecated)]
pub use cast::cast;
pub use compare::*;
pub use fill_null::*;
pub use filter::*;
#[expect(deprecated)]
pub use invert::invert;
pub use is_constant::*;
pub use is_sorted::*;
use itertools::Itertools;
pub use list_contains::*;
pub use mask::*;
pub use min_max::*;
pub use nan_count::*;
pub use numeric::*;
use parking_lot::RwLock;
pub use sum::*;
use vortex_dtype::DType;
use vortex_error::VortexError;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_err;
use vortex_mask::Mask;
pub use zip::*;

use crate::Array;
use crate::ArrayRef;
use crate::builders::ArrayBuilder;
pub use crate::expr::BetweenExecuteAdaptor;
pub use crate::expr::BetweenKernel;
pub use crate::expr::BetweenReduce;
pub use crate::expr::BetweenReduceAdaptor;
pub use crate::expr::CastExecuteAdaptor;
pub use crate::expr::CastKernel;
pub use crate::expr::CastReduce;
pub use crate::expr::CastReduceAdaptor;
pub use crate::expr::FillNullExecuteAdaptor;
pub use crate::expr::FillNullKernel;
pub use crate::expr::FillNullReduce;
pub use crate::expr::FillNullReduceAdaptor;
pub use crate::expr::MaskExecuteAdaptor;
pub use crate::expr::MaskKernel;
pub use crate::expr::MaskReduce;
pub use crate::expr::MaskReduceAdaptor;
pub use crate::expr::NotExecuteAdaptor;
pub use crate::expr::NotKernel;
pub use crate::expr::NotReduce;
pub use crate::expr::NotReduceAdaptor;
use crate::scalar::Scalar;

#[cfg(feature = "arbitrary")]
mod arbitrary;
mod boolean;
mod cast;
mod compare;
#[cfg(feature = "_test-harness")]
pub mod conformance;
mod fill_null;
mod filter;
mod invert;
mod is_constant;
mod is_sorted;
mod list_contains;
mod mask;
mod min_max;
mod nan_count;
mod numeric;
mod sum;
mod zip;

/// An instance of a compute function holding the implementation vtable and a set of registered
/// compute kernels.
pub struct ComputeFn {
    id: ArcRef<str>,
    vtable: ArcRef<dyn ComputeFnVTable>,
    kernels: RwLock<Vec<ArcRef<dyn Kernel>>>,
}

/// Force all the default [`ComputeFn`] vtables to register all available compute kernels.
///
/// Mostly useful for small benchmarks where the overhead might cause noise depending on the order of benchmarks.
pub fn warm_up_vtables() {
    #[allow(unused_qualifications)]
    is_constant::warm_up_vtable();
    is_sorted::warm_up_vtable();
    list_contains::warm_up_vtable();
    min_max::warm_up_vtable();
    nan_count::warm_up_vtable();
    sum::warm_up_vtable();
}

impl ComputeFn {
    /// Create a new compute function from the given [`ComputeFnVTable`].
    pub fn new(id: ArcRef<str>, vtable: ArcRef<dyn ComputeFnVTable>) -> Self {
        Self {
            id,
            vtable,
            kernels: Default::default(),
        }
    }

    /// Returns the string identifier of the compute function.
    pub fn id(&self) -> &ArcRef<str> {
        &self.id
    }

    /// Register a kernel for the compute function.
    pub fn register_kernel(&self, kernel: ArcRef<dyn Kernel>) {
        self.kernels.write().push(kernel);
    }

    /// Invokes the compute function with the given arguments.
    pub fn invoke(&self, args: &InvocationArgs) -> VortexResult<Output> {
        // Perform some pre-condition checks against the arguments and the function properties.
        if self.is_elementwise() {
            // For element-wise functions, all input arrays must be the same length.
            if !args
                .inputs
                .iter()
                .filter_map(|input| input.array())
                .map(|array| array.len())
                .all_equal()
            {
                vortex_bail!(
                    "Compute function {} is elementwise but input arrays have different lengths",
                    self.id
                );
            }
        }

        let expected_dtype = self.vtable.return_dtype(args)?;
        let expected_len = self.vtable.return_len(args)?;

        let output = self.vtable.invoke(args, &self.kernels.read())?;

        if output.dtype() != &expected_dtype {
            vortex_bail!(
                "Internal error: compute function {} returned a result of type {} but expected {}\n{}",
                self.id,
                output.dtype(),
                &expected_dtype,
                args.inputs
                    .iter()
                    .filter_map(|input| input.array())
                    .format_with(",", |array, f| f(&array.encoding_id()))
            );
        }
        if output.len() != expected_len {
            vortex_bail!(
                "Internal error: compute function {} returned a result of length {} but expected {}",
                self.id,
                output.len(),
                expected_len
            );
        }

        Ok(output)
    }

    /// Compute the return type of the function given the input arguments.
    pub fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        self.vtable.return_dtype(args)
    }

    /// Compute the return length of the function given the input arguments.
    pub fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        self.vtable.return_len(args)
    }

    /// Returns whether the compute function is elementwise, i.e. the output is the same shape as
    pub fn is_elementwise(&self) -> bool {
        // TODO(ngates): should this just be a constant passed in the constructor?
        self.vtable.is_elementwise()
    }

    /// Returns the compute function's kernels.
    pub fn kernels(&self) -> Vec<ArcRef<dyn Kernel>> {
        self.kernels.read().to_vec()
    }
}

/// VTable for the implementation of a compute function.
pub trait ComputeFnVTable: 'static + Send + Sync {
    /// Invokes the compute function entry-point with the given input arguments and options.
    ///
    /// The entry-point logic can short-circuit compute using statistics, update result array
    /// statistics, search for relevant compute kernels, and canonicalize the inputs in order
    /// to successfully compute a result.
    fn invoke(&self, args: &InvocationArgs, kernels: &[ArcRef<dyn Kernel>])
    -> VortexResult<Output>;

    /// Computes the return type of the function given the input arguments.
    ///
    /// All kernel implementations will be validated to return the [`DType`] as computed here.
    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType>;

    /// Computes the return length of the function given the input arguments.
    ///
    /// All kernel implementations will be validated to return the len as computed here.
    /// Scalars are considered to have length 1.
    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize>;

    /// Returns whether the function operates elementwise, i.e. the output is the same shape as the
    /// input and no information is shared between elements.
    ///
    /// Examples include `add`, `subtract`, `and`, `cast`, `fill_null` etc.
    /// Examples that are not elementwise include `sum`, `count`, `min`, `fill_forward` etc.
    ///
    /// All input arrays to an elementwise function *must* have the same length.
    fn is_elementwise(&self) -> bool;
}

/// Arguments to a compute function invocation.
#[derive(Clone)]
pub struct InvocationArgs<'a> {
    pub inputs: &'a [Input<'a>],
    pub options: &'a dyn Options,
}

/// For unary compute functions, it's useful to just have this short-cut.
pub struct UnaryArgs<'a, O: Options> {
    pub array: &'a dyn Array,
    pub options: &'a O,
}

impl<'a, O: Options> TryFrom<&InvocationArgs<'a>> for UnaryArgs<'a, O> {
    type Error = VortexError;

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 1 {
            vortex_bail!("Expected 1 input, found {}", value.inputs.len());
        }
        let array = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
        let options =
            value.options.as_any().downcast_ref::<O>().ok_or_else(|| {
                vortex_err!("Expected options to be of type {}", type_name::<O>())
            })?;
        Ok(UnaryArgs { array, options })
    }
}

/// For binary compute functions, it's useful to just have this short-cut.
pub struct BinaryArgs<'a, O: Options> {
    pub lhs: &'a dyn Array,
    pub rhs: &'a dyn Array,
    pub options: &'a O,
}

impl<'a, O: Options> TryFrom<&InvocationArgs<'a>> for BinaryArgs<'a, O> {
    type Error = VortexError;

    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
        if value.inputs.len() != 2 {
            vortex_bail!("Expected 2 input, found {}", value.inputs.len());
        }
        let lhs = value.inputs[0]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
        let rhs = value.inputs[1]
            .array()
            .ok_or_else(|| vortex_err!("Expected input 1 to be an array"))?;
        let options =
            value.options.as_any().downcast_ref::<O>().ok_or_else(|| {
                vortex_err!("Expected options to be of type {}", type_name::<O>())
            })?;
        Ok(BinaryArgs { lhs, rhs, options })
    }
}

/// Input to a compute function.
pub enum Input<'a> {
    Scalar(&'a Scalar),
    Array(&'a dyn Array),
    Mask(&'a Mask),
    Builder(&'a mut dyn ArrayBuilder),
    DType(&'a DType),
}

impl Debug for Input<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let mut f = f.debug_struct("Input");
        match self {
            Input::Scalar(scalar) => f.field("Scalar", scalar),
            Input::Array(array) => f.field("Array", array),
            Input::Mask(mask) => f.field("Mask", mask),
            Input::Builder(builder) => f.field("Builder", &builder.len()),
            Input::DType(dtype) => f.field("DType", dtype),
        };
        f.finish()
    }
}

impl<'a> From<&'a dyn Array> for Input<'a> {
    fn from(value: &'a dyn Array) -> Self {
        Input::Array(value)
    }
}

impl<'a> From<&'a Scalar> for Input<'a> {
    fn from(value: &'a Scalar) -> Self {
        Input::Scalar(value)
    }
}

impl<'a> From<&'a Mask> for Input<'a> {
    fn from(value: &'a Mask) -> Self {
        Input::Mask(value)
    }
}

impl<'a> From<&'a DType> for Input<'a> {
    fn from(value: &'a DType) -> Self {
        Input::DType(value)
    }
}

impl<'a> Input<'a> {
    pub fn scalar(&self) -> Option<&'a Scalar> {
        if let Input::Scalar(scalar) = self {
            Some(*scalar)
        } else {
            None
        }
    }

    pub fn array(&self) -> Option<&'a dyn Array> {
        if let Input::Array(array) = self {
            Some(*array)
        } else {
            None
        }
    }

    pub fn mask(&self) -> Option<&'a Mask> {
        if let Input::Mask(mask) = self {
            Some(*mask)
        } else {
            None
        }
    }

    pub fn builder(&'a mut self) -> Option<&'a mut dyn ArrayBuilder> {
        if let Input::Builder(builder) = self {
            Some(*builder)
        } else {
            None
        }
    }

    pub fn dtype(&self) -> Option<&'a DType> {
        if let Input::DType(dtype) = self {
            Some(*dtype)
        } else {
            None
        }
    }
}

/// Output from a compute function.
#[derive(Debug)]
pub enum Output {
    Scalar(Scalar),
    Array(ArrayRef),
}

#[expect(
    clippy::len_without_is_empty,
    reason = "Output is always non-empty (scalar has len 1)"
)]
impl Output {
    pub fn dtype(&self) -> &DType {
        match self {
            Output::Scalar(scalar) => scalar.dtype(),
            Output::Array(array) => array.dtype(),
        }
    }

    pub fn len(&self) -> usize {
        match self {
            Output::Scalar(_) => 1,
            Output::Array(array) => array.len(),
        }
    }

    pub fn unwrap_scalar(self) -> VortexResult<Scalar> {
        match self {
            Output::Array(_) => vortex_bail!("Expected scalar output, got Array"),
            Output::Scalar(scalar) => Ok(scalar),
        }
    }

    pub fn unwrap_array(self) -> VortexResult<ArrayRef> {
        match self {
            Output::Array(array) => Ok(array),
            Output::Scalar(_) => vortex_bail!("Expected array output, got Scalar"),
        }
    }
}

impl From<ArrayRef> for Output {
    fn from(value: ArrayRef) -> Self {
        Output::Array(value)
    }
}

impl From<Scalar> for Output {
    fn from(value: Scalar) -> Self {
        Output::Scalar(value)
    }
}

/// Options for a compute function invocation.
pub trait Options: 'static {
    fn as_any(&self) -> &dyn Any;
}

impl Options for () {
    fn as_any(&self) -> &dyn Any {
        self
    }
}

/// Compute functions can ask arrays for compute kernels for a given invocation.
///
/// The kernel is invoked with the input arguments and options, and can return `None` if it is
/// unable to compute the result for the given inputs due to missing implementation logic.
/// For example, if kernel doesn't support the `LTE` operator. By returning `None`, the kernel
/// is indicating that it cannot compute the result for the given inputs, and another kernel should
/// be tried. *Not* that the given inputs are invalid for the compute function.
///
/// If the kernel fails to compute a result, it should return a `Some` with the error.
pub trait Kernel: 'static + Send + Sync + Debug {
    /// Invokes the kernel with the given input arguments and options.
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>>;
}

/// Register a kernel for a compute function.
/// See each compute function for the correct type of kernel to register.
#[macro_export]
macro_rules! register_kernel {
    ($T:expr) => {
        $crate::aliases::inventory::submit!($T);
    };
}