vortex_array/arrays/list/compute/

take.rs

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

use vortex_dtype::IntegerPType;
use vortex_dtype::Nullability;
use vortex_dtype::match_each_integer_ptype;
use vortex_dtype::match_smallest_offset_type;
use vortex_error::VortexExpect;
use vortex_error::VortexResult;

use crate::Array;
use crate::ArrayRef;
use crate::ToCanonical;
use crate::arrays::ListArray;
use crate::arrays::ListVTable;
use crate::arrays::PrimitiveArray;
use crate::builders::ArrayBuilder;
use crate::builders::PrimitiveBuilder;
use crate::compute::TakeKernel;
use crate::compute::TakeKernelAdapter;
use crate::compute::take;
use crate::register_kernel;
use crate::vtable::ValidityHelper;

// TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call
// the `ListView::take` compute function once `ListView` is more stable.

/// Take implementation for [`ListArray`].
///
/// Unlike `ListView`, `ListArray` must rebuild the elements array to maintain its invariant
/// that lists are stored contiguously and in-order (`offset[i+1] >= offset[i]`). Taking
/// non-contiguous indices would violate this requirement.
impl TakeKernel for ListVTable {
    #[expect(clippy::cognitive_complexity)]
    fn take(&self, array: &ListArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
        let indices = indices.to_primitive();
        // This is an over-approximation of the total number of elements in the resulting array.
        let total_approx = array.elements().len().saturating_mul(indices.len());

        match_each_integer_ptype!(array.offsets().dtype().as_ptype(), |O| {
            match_each_integer_ptype!(indices.ptype(), |I| {
                match_smallest_offset_type!(total_approx, |OutputOffsetType| {
                    _take::<I, O, OutputOffsetType>(array, &indices)
                })
            })
        })
    }
}

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

fn _take<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
    array: &ListArray,
    indices_array: &PrimitiveArray,
) -> VortexResult<ArrayRef> {
    let data_validity = array.validity_mask();
    let indices_validity = indices_array.validity_mask();

    if !indices_validity.all_true() || !data_validity.all_true() {
        return _take_nullable::<I, O, OutputOffsetType>(array, indices_array);
    }

    let offsets_array = array.offsets().to_primitive();
    let offsets: &[O] = offsets_array.as_slice();
    let indices: &[I] = indices_array.as_slice();

    let mut new_offsets = PrimitiveBuilder::<OutputOffsetType>::with_capacity(
        Nullability::NonNullable,
        indices.len(),
    );
    let mut elements_to_take =
        PrimitiveBuilder::with_capacity(Nullability::NonNullable, 2 * indices.len());

    let mut current_offset = OutputOffsetType::zero();
    new_offsets.append_zero();

    for &data_idx in indices {
        let data_idx: usize = data_idx.as_();

        let start = offsets[data_idx];
        let stop = offsets[data_idx + 1];

        // Annoyingly, we can't turn (start..end) into a range, so we're doing that manually.
        //
        // We could convert start and end to usize, but that would impose a potentially
        // harder constraint - now we don't care if they fit into usize as long as their
        // difference does.
        let additional: usize = (stop - start).as_();

        // TODO(0ax1): optimize this
        elements_to_take.reserve_exact(additional);
        for i in 0..additional {
            elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional"));
        }
        current_offset +=
            OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
        new_offsets.append_value(current_offset);
    }

    let elements_to_take = elements_to_take.finish();
    let new_offsets = new_offsets.finish();

    let new_elements = take(array.elements(), elements_to_take.as_ref())?;

    Ok(ListArray::try_new(
        new_elements,
        new_offsets,
        array.validity().clone().take(indices_array.as_ref())?,
    )?
    .to_array())
}

fn _take_nullable<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
    array: &ListArray,
    indices_array: &PrimitiveArray,
) -> VortexResult<ArrayRef> {
    let offsets_array = array.offsets().to_primitive();
    let offsets: &[O] = offsets_array.as_slice();
    let indices: &[I] = indices_array.as_slice();
    let data_validity = array.validity_mask();
    let indices_validity = indices_array.validity_mask();

    let mut new_offsets = PrimitiveBuilder::<OutputOffsetType>::with_capacity(
        Nullability::NonNullable,
        indices.len(),
    );

    // This will be the indices we push down to the child array to call `take` with.
    //
    // There are 2 things to note here:
    // - We do not know how many elements we need to take from our child since lists are variable
    //   size: thus we arbitrarily choose a capacity of `2 * # of indices`.
    // - The type of the primitive builder needs to fit the largest offset of the (parent)
    //   `ListArray`, so we make this `PrimitiveBuilder` generic over `O` (instead of `I`).
    let mut elements_to_take =
        PrimitiveBuilder::<O>::with_capacity(Nullability::NonNullable, 2 * indices.len());

    let mut current_offset = OutputOffsetType::zero();
    new_offsets.append_zero();

    for (idx, data_idx) in indices.iter().enumerate() {
        if !indices_validity.value(idx) {
            new_offsets.append_value(current_offset);
            continue;
        }

        let data_idx: usize = data_idx.as_();

        if !data_validity.value(data_idx) {
            new_offsets.append_value(current_offset);
            continue;
        }

        let start = offsets[data_idx];
        let stop = offsets[data_idx + 1];

        // See the note in `_take` on the reasoning.
        let additional: usize = (stop - start).as_();

        elements_to_take.reserve_exact(additional);
        for i in 0..additional {
            elements_to_take.append_value(start + O::from_usize(i).vortex_expect("i < additional"));
        }
        current_offset +=
            OutputOffsetType::from_usize((stop - start).as_()).vortex_expect("offset conversion");
        new_offsets.append_value(current_offset);
    }

    let elements_to_take = elements_to_take.finish();
    let new_offsets = new_offsets.finish();
    let new_elements = take(array.elements(), elements_to_take.as_ref())?;

    Ok(ListArray::try_new(
        new_elements,
        new_offsets,
        array.validity().clone().take(indices_array.as_ref())?,
    )?
    .to_array())
}

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

    use rstest::rstest;
    use vortex_buffer::buffer;
    use vortex_dtype::DType;
    use vortex_dtype::Nullability;
    use vortex_dtype::PType::I32;
    use vortex_scalar::Scalar;

    use crate::Array;
    use crate::IntoArray as _;
    use crate::ToCanonical;
    use crate::arrays::BoolArray;
    use crate::arrays::PrimitiveArray;
    use crate::arrays::list::ListArray;
    use crate::compute::conformance::take::test_take_conformance;
    use crate::compute::take;
    use crate::validity::Validity;

    #[test]
    fn nullable_take() {
        let list = ListArray::try_new(
            buffer![0i32, 5, 3, 4].into_array(),
            buffer![0, 2, 3, 4, 4].into_array(),
            Validity::Array(BoolArray::from_iter(vec![true, true, false, true]).to_array()),
        )
        .unwrap()
        .to_array();

        let idx =
            PrimitiveArray::from_option_iter(vec![Some(0), None, Some(1), Some(3)]).to_array();

        let result = take(&list, &idx).unwrap();

        assert_eq!(
            result.dtype(),
            &DType::List(
                Arc::new(DType::Primitive(I32, Nullability::NonNullable)),
                Nullability::Nullable
            )
        );

        let result = result.to_listview();

        assert_eq!(result.len(), 4);

        let element_dtype: Arc<DType> = Arc::new(I32.into());

        assert!(result.is_valid(0));
        assert_eq!(
            result.scalar_at(0),
            Scalar::list(
                element_dtype.clone(),
                vec![0i32.into(), 5.into()],
                Nullability::Nullable
            )
        );

        assert!(result.is_invalid(1));

        assert!(result.is_valid(2));
        assert_eq!(
            result.scalar_at(2),
            Scalar::list(
                element_dtype.clone(),
                vec![3i32.into()],
                Nullability::Nullable
            )
        );

        assert!(result.is_valid(3));
        assert_eq!(
            result.scalar_at(3),
            Scalar::list(element_dtype, vec![], Nullability::Nullable)
        );
    }

    #[test]
    fn change_validity() {
        let list = ListArray::try_new(
            buffer![0i32, 5, 3, 4].into_array(),
            buffer![0, 2, 3].into_array(),
            Validity::NonNullable,
        )
        .unwrap()
        .to_array();

        let idx = PrimitiveArray::from_option_iter(vec![Some(0), Some(1), None]).to_array();
        // since idx is nullable, the final list will also be nullable

        let result = take(&list, &idx).unwrap();
        assert_eq!(
            result.dtype(),
            &DType::List(
                Arc::new(DType::Primitive(I32, Nullability::NonNullable)),
                Nullability::Nullable
            )
        );
    }

    #[test]
    fn non_nullable_take() {
        let list = ListArray::try_new(
            buffer![0i32, 5, 3, 4].into_array(),
            buffer![0, 2, 3, 3, 4].into_array(),
            Validity::NonNullable,
        )
        .unwrap()
        .to_array();

        let idx = buffer![1, 0, 2].into_array();

        let result = take(&list, &idx).unwrap();

        assert_eq!(
            result.dtype(),
            &DType::List(
                Arc::new(DType::Primitive(I32, Nullability::NonNullable)),
                Nullability::NonNullable
            )
        );

        let result = result.to_listview();

        assert_eq!(result.len(), 3);

        let element_dtype: Arc<DType> = Arc::new(I32.into());

        assert!(result.is_valid(0));
        assert_eq!(
            result.scalar_at(0),
            Scalar::list(
                element_dtype.clone(),
                vec![3i32.into()],
                Nullability::NonNullable
            )
        );

        assert!(result.is_valid(1));
        assert_eq!(
            result.scalar_at(1),
            Scalar::list(
                element_dtype.clone(),
                vec![0i32.into(), 5.into()],
                Nullability::NonNullable
            )
        );

        assert!(result.is_valid(2));
        assert_eq!(
            result.scalar_at(2),
            Scalar::list(element_dtype, vec![], Nullability::NonNullable)
        );
    }

    #[test]
    fn test_take_empty_array() {
        let list = ListArray::try_new(
            buffer![0i32, 5, 3, 4].into_array(),
            buffer![0].into_array(),
            Validity::NonNullable,
        )
        .unwrap()
        .to_array();

        let idx = PrimitiveArray::empty::<i32>(Nullability::Nullable).to_array();

        let result = take(&list, &idx).unwrap();
        assert_eq!(
            result.dtype(),
            &DType::List(
                Arc::new(DType::Primitive(I32, Nullability::NonNullable)),
                Nullability::Nullable
            )
        );
        assert_eq!(result.len(), 0,);
    }

    #[rstest]
    #[case(ListArray::try_new(
        buffer![0i32, 1, 2, 3, 4, 5].into_array(),
        buffer![0, 2, 3, 5, 5, 6].into_array(),
        Validity::NonNullable,
    ).unwrap())]
    #[case(ListArray::try_new(
        buffer![10i32, 20, 30, 40, 50].into_array(),
        buffer![0, 2, 3, 4, 5].into_array(),
        Validity::Array(BoolArray::from_iter(vec![true, false, true, true]).to_array()),
    ).unwrap())]
    #[case(ListArray::try_new(
        buffer![1i32, 2, 3].into_array(),
        buffer![0, 0, 2, 2, 3].into_array(), // First and third are empty
        Validity::NonNullable,
    ).unwrap())]
    #[case(ListArray::try_new(
        buffer![42i32, 43].into_array(),
        buffer![0, 2].into_array(),
        Validity::NonNullable,
    ).unwrap())]
    #[case({
        let elements = buffer![0i32..200].into_array();
        let mut offsets = vec![0u64];
        for i in 1..=50 {
            offsets.push(offsets[i - 1] + (i as u64 % 5)); // Variable length lists
        }
        ListArray::try_new(
            elements,
            PrimitiveArray::from_iter(offsets).to_array(),
            Validity::NonNullable,
        ).unwrap()
    })]
    #[case(ListArray::try_new(
        PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]).to_array(),
        buffer![0, 2, 3, 5].into_array(),
        Validity::NonNullable,
    ).unwrap())]
    fn test_take_list_conformance(#[case] list: ListArray) {
        test_take_conformance(list.as_ref());
    }

    #[test]
    fn test_u64_offset_accumulation_non_nullable() {
        let elements = buffer![0i32; 200].into_array();
        let offsets = buffer![0u8, 200].into_array();
        let list = ListArray::try_new(elements, offsets, Validity::NonNullable)
            .unwrap()
            .to_array();

        // Take the same large list twice - would overflow u8 but works with u64.
        let idx = buffer![0u8, 0].into_array();
        let result = take(&list, &idx).unwrap();

        assert_eq!(result.len(), 2);

        let result_view = result.to_listview();
        assert_eq!(result_view.len(), 2);
        assert!(result_view.is_valid(0));
        assert!(result_view.is_valid(1));
    }

    #[test]
    fn test_u64_offset_accumulation_nullable() {
        let elements = buffer![0i32; 150].into_array();
        let offsets = buffer![0u8, 150, 150].into_array();
        let validity = BoolArray::from_iter(vec![true, false]).to_array();
        let list = ListArray::try_new(elements, offsets, Validity::Array(validity))
            .unwrap()
            .to_array();

        // Take the same large list twice - would overflow u8 but works with u64.
        let idx = PrimitiveArray::from_option_iter(vec![Some(0u8), None, Some(0u8)]).to_array();
        let result = take(&list, &idx).unwrap();

        assert_eq!(result.len(), 3);

        let result_view = result.to_listview();
        assert_eq!(result_view.len(), 3);
        assert!(result_view.is_valid(0));
        assert!(result_view.is_invalid(1));
        assert!(result_view.is_valid(2));
    }

    /// Regression test for validity length mismatch bug.
    ///
    /// When source array has `Validity::Array(...)` and indices are non-nullable,
    /// the result validity must have length equal to indices.len(), not source.len().
    #[test]
    fn test_take_validity_length_mismatch_regression() {
        // Source array with explicit validity array (length 2).
        let list = ListArray::try_new(
            buffer![1i32, 2, 3, 4].into_array(),
            buffer![0, 2, 4].into_array(),
            Validity::Array(BoolArray::from_iter(vec![true, true]).to_array()),
        )
        .unwrap()
        .to_array();

        // Take more indices than source length (4 vs 2) with non-nullable indices.
        let idx = buffer![0u32, 1, 0, 1].into_array();

        // This should not panic - result should have length 4.
        let result = take(&list, &idx).unwrap();
        assert_eq!(result.len(), 4);
    }
}