vortex-array 0.88.0

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

//! Execution for primitive [`Interleave`] values.

use num_traits::AsPrimitive;
use vortex_buffer::Buffer;
use vortex_buffer::BufferAllocatorRef;
use vortex_buffer::BufferMut;
use vortex_error::VortexResult;
use vortex_error::vortex_bail;
use vortex_error::vortex_ensure;
use vortex_error::vortex_ensure_eq;

use super::super::Interleave;
use super::super::InterleaveArrayExt;
use crate::AnyColumnar;
use crate::array::Array;
use crate::arrays::Constant;
use crate::arrays::Primitive;
use crate::arrays::PrimitiveArray;
use crate::arrays::primitive::PrimitiveArrayExt;
use crate::dtype::NativePType;
use crate::executor::ExecutionCtx;
use crate::executor::ExecutionResult;
use crate::match_each_native_ptype;
use crate::match_each_unsigned_integer_ptype;
use crate::require_child;

pub(super) fn execute(
    mut array: Array<Interleave>,
    ctx: &mut ExecutionCtx,
) -> VortexResult<ExecutionResult> {
    let num_values = array.num_values();
    array = require_child!(array, array.array_indices(), 0 => Primitive);
    array = require_child!(array, array.row_indices(), 1 => Primitive);
    for i in 0..num_values {
        array = require_child!(array, array.value(i), i + 2 => AnyColumnar);
    }

    let validity = array.as_ref().validity()?;
    let output = match_each_native_ptype!(array.dtype().as_ptype(), |T| {
        let values = gather_values::<T>(&array, ctx.allocator())?;
        VortexResult::Ok(PrimitiveArray::new(values, validity))
    })?;

    Ok(ExecutionResult::done(output))
}

/// Physical storage for one primitive source array.
struct PrimitiveSource<T> {
    data: Buffer<T>,
    len: usize,
    // Maps logical rows to `data`: zero for constants and identity for decoded buffers.
    row_mask: usize,
}

fn gather_values<T: NativePType>(
    array: &Array<Interleave>,
    allocator: &BufferAllocatorRef,
) -> VortexResult<Buffer<T>> {
    let values = (0..array.num_values())
        .map(|i| {
            let value = array.value(i);
            let len = value.len();
            if let Some(constant) = value.as_opt::<Constant>() {
                // Validity carries nullness; a null constant's payload is never observed.
                let payload = constant
                    .scalar()
                    .as_primitive()
                    .typed_value::<T>()
                    .unwrap_or_default();
                PrimitiveSource {
                    data: Buffer::full_in(payload, 1, allocator.clone()),
                    len,
                    row_mask: 0,
                }
            } else {
                PrimitiveSource {
                    data: value.as_::<Primitive>().to_buffer::<T>(),
                    len,
                    row_mask: usize::MAX,
                }
            }
        })
        .collect::<Vec<_>>();
    let branches = array.array_indices().as_::<Primitive>();
    let rows = array.row_indices().as_::<Primitive>();

    match_each_unsigned_integer_ptype!(branches.ptype(), |A| {
        match_each_unsigned_integer_ptype!(rows.ptype(), |R| {
            gather(
                &values,
                branches.as_slice::<A>(),
                rows.as_slice::<R>(),
                allocator,
            )
        })
    })
}

fn gather<T, A, R>(
    values: &[PrimitiveSource<T>],
    branches: &[A],
    rows: &[R],
    allocator: &BufferAllocatorRef,
) -> VortexResult<Buffer<T>>
where
    T: NativePType,
    A: AsPrimitive<usize>,
    R: AsPrimitive<usize>,
{
    // `zip` truncates to the shorter input.
    vortex_ensure_eq!(
        rows.len(),
        branches.len(),
        "interleave row_indices length does not match array_indices length",
    );

    let mut output = BufferMut::with_capacity_in(branches.len(), allocator.clone());
    output.try_extend_trusted(branches.iter().zip(rows).map(|(branch, row)| {
        let Some(source) = values.get((*branch).as_()) else {
            vortex_bail!("interleave array index out of bounds");
        };
        let row = (*row).as_();
        vortex_ensure!(row < source.len, "interleave row index out of bounds");
        Ok(source.data[row & source.row_mask])
    }))?;
    Ok(output.freeze())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::memory::BufferAllocatorRef;

    #[test]
    fn rejects_out_of_bounds_selectors() {
        let values = [PrimitiveSource {
            data: Buffer::full_in(1u32, 1, BufferAllocatorRef::statically_allocated()),
            len: 1,
            row_mask: 0,
        }];

        let allocator = BufferAllocatorRef::statically_allocated();
        assert!(gather(&values, &[1u8], &[0u8], &allocator).is_err());
        assert!(gather(&values, &[0u8], &[1u8], &allocator).is_err());
    }
}