16331938722

Committed 16 Jul 2025 10:49PM UTC coverage: 80.702% (-0.9%) from 81.557%

Build # 16331938722

Build Type

push

github

Committed by

web-flow

Commit Message

feat: build with stable rust (#3881)

Run Details

120 of 173 new or added lines in 28 files covered. (69.36%)

174 existing lines in 102 files now uncovered.

41861 of 51871 relevant lines covered (80.7%)

157487.71 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

81.29

/vortex-array/src/compute/list_contains.rs

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

//! List-related compute operations.

use std::sync::LazyLock;

use arcref::ArcRef;
use arrow_buffer::BooleanBuffer;
use arrow_buffer::bit_iterator::BitIndexIterator;
use num_traits::AsPrimitive;
use vortex_buffer::Buffer;
use vortex_dtype::{DType, NativePType, Nullability, match_each_integer_ptype};
use vortex_error::{VortexExpect, VortexResult, vortex_bail};
use vortex_scalar::{ListScalar, Scalar};

use crate::arrays::{BoolArray, ConstantArray, ListArray};
use crate::compute::{
    BinaryArgs, ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Operator, Output, compare,
    fill_null, or,
};
use crate::validity::Validity;
use crate::vtable::{VTable, ValidityHelper};
use crate::{Array, ArrayRef, IntoArray, ToCanonical};

/// Compute a `Bool`-typed array the same length as `array` where elements is `true` if the list
/// item contains the `value`, `false` otherwise.
///
/// ## Null scalar handling
///
/// If the `value` or `array` is `null` at any index the result at that index is `null`.
///
/// ## Format semantics
/// ```txt
/// list_contains(list, elem)
///   ==> (!is_null(list) or NULL) and (!is_null(elem) or NULL) and any({elem = elem_i | elem_i in list}),
/// ```
///
/// ## Example
///
/// ```rust
/// use vortex_array::{Array, IntoArray, ToCanonical};
/// use vortex_array::arrays::{ConstantArray, ListArray, VarBinArray};
/// use vortex_array::compute::list_contains;
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_dtype::DType;
/// use vortex_scalar::Scalar;
/// let elements = VarBinArray::from_vec(
///         vec!["a", "a", "b", "a", "c"], DType::Utf8(false.into())).into_array();
/// let offsets = buffer![0u32, 1, 3, 5].into_array();
/// let list_array = ListArray::try_new(elements, offsets, Validity::NonNullable).unwrap();
///
/// let matches = list_contains(list_array.as_ref(), ConstantArray::new(Scalar::from("b"), list_array.len()).as_ref()).unwrap();
/// let to_vec: Vec<bool> = matches.to_bool().unwrap().boolean_buffer().iter().collect();
/// assert_eq!(to_vec, vec![false, true, false]);
/// ```
pub fn list_contains(array: &dyn Array, value: &dyn Array) -> VortexResult<ArrayRef> {
    LIST_CONTAINS_FN
        .invoke(&InvocationArgs {
            inputs: &[array.into(), value.into()],
            options: &(),
        })?
        .unwrap_array()
}

pub struct ListContains;

impl ComputeFnVTable for ListContains {
    fn invoke(
        &self,
        args: &InvocationArgs,
        kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        let BinaryArgs {
            lhs: array,
            rhs: value,
            ..
        } = BinaryArgs::<()>::try_from(args)?;

        let DType::List(elem_dtype, _) = array.dtype() else {
            vortex_bail!("Array must be of List type");
        };
        if !elem_dtype.as_ref().eq_ignore_nullability(value.dtype()) {
            vortex_bail!(
                "Element type {} of ListArray does not match search value {}",
                elem_dtype,
                value.dtype(),
            );
        };

        if value.all_invalid()? || array.all_invalid()? {
            return Ok(Output::Array(
                ConstantArray::new(
                    Scalar::null(DType::Bool(Nullability::Nullable)),
                    array.len(),
                )
                .to_array(),
            ));
        }

        for kernel in kernels {
            if let Some(output) = kernel.invoke(args)? {
                return Ok(output);
            }
        }
        if let Some(output) = array.invoke(&LIST_CONTAINS_FN, args)? {
            return Ok(output);
        }

        let nullability = array.dtype().nullability() | value.dtype().nullability();

        let result = if let Some(value_scalar) = value.as_constant() {
            list_contains_scalar(array, &value_scalar, nullability)
        } else if let Some(list_scalar) = array.as_constant() {
            constant_list_scalar_contains(&list_scalar.as_list(), value, nullability)
        } else {
            todo!("unsupported list contains with list and element as arrays")
        };

        result.map(Output::Array)
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        let input = BinaryArgs::<()>::try_from(args)?;
        Ok(DType::Bool(
            input.lhs.dtype().nullability() | input.rhs.dtype().nullability(),
        ))
    }

    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
        Ok(BinaryArgs::<()>::try_from(args)?.lhs.len())
    }

    fn is_elementwise(&self) -> bool {
        true
    }
}

pub trait ListContainsKernel: VTable {
    fn list_contains(
        &self,
        list: &dyn Array,
        element: &Self::Array,
    ) -> VortexResult<Option<ArrayRef>>;
}

pub struct ListContainsKernelRef(ArcRef<dyn Kernel>);
inventory::collect!(ListContainsKernelRef);

#[derive(Debug)]
pub struct ListContainsKernelAdapter<V: VTable>(pub V);

impl<V: VTable + ListContainsKernel> ListContainsKernelAdapter<V> {
    pub const fn lift(&'static self) -> ListContainsKernelRef {
        ListContainsKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + ListContainsKernel> Kernel for ListContainsKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let BinaryArgs {
            lhs: array,
            rhs: value,
            ..
        } = BinaryArgs::<()>::try_from(args)?;
        let Some(value) = value.as_opt::<V>() else {
            return Ok(None);
        };
        self.0
            .list_contains(array, value)
            .map(|c| c.map(Output::Array))
    }
}

pub static LIST_CONTAINS_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("list_contains".into(), ArcRef::new_ref(&ListContains));
    for kernel in inventory::iter::<ListContainsKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

// Then there is a constant list scalar (haystack) being compared to an array of needles.
fn constant_list_scalar_contains(
    list_scalar: &ListScalar,
    values: &dyn Array,
    nullability: Nullability,
) -> VortexResult<ArrayRef> {
    let elements = list_scalar.elements().vortex_expect("non null");

    let len = values.len();
    let mut result: Option<ArrayRef> = None;
    let false_scalar = Scalar::bool(false, nullability);
    for element in elements {
        let res = fill_null(
            &compare(
                ConstantArray::new(element, len).as_ref(),
                values,
                Operator::Eq,
            )?,
            &false_scalar,
        )?;
        if let Some(acc) = result {
            result = Some(or(&acc, &res)?)
        } else {
            result = Some(res);
        }
    }
    Ok(result.unwrap_or_else(|| ConstantArray::new(false_scalar, len).to_array()))
}

fn list_contains_scalar(
    array: &dyn Array,
    value: &Scalar,
    nullability: Nullability,
) -> VortexResult<ArrayRef> {
    // If the list array is constant, we perform a single comparison.
    if array.len() > 1 && array.is_constant() {
        let contains = list_contains_scalar(&array.slice(0, 1)?, value, nullability)?;
        return Ok(ConstantArray::new(contains.scalar_at(0)?, array.len()).into_array());
    }

    // Canonicalize to a list array.
    // NOTE(ngates): we may wish to add elements and offsets accessors to the ListArrayTrait.
    let list_array = array.to_list()?;

    let elems = list_array.elements();
    if elems.is_empty() {
        // Must return false when a list is empty (but valid), or null when the list itself is null.
        return list_false_or_null(&list_array, nullability);
    }

    let rhs = ConstantArray::new(value.clone(), elems.len());
    let matching_elements = compare(elems, rhs.as_ref(), Operator::Eq)?;
    let matches = matching_elements.to_bool()?;

    // Fast path: no elements match.
    if let Some(pred) = matches.as_constant() {
        return match pred.as_bool().value() {
            // All comparisons are invalid (result in `null`), and search is not null because
            // we already checked for null above.
            None => {
                assert!(
                    !rhs.scalar().is_null(),
                    "Search value must not be null here"
                );
                // False, unless the list itself is null in which case we return null.
                list_false_or_null(&list_array, nullability)
            }
            // No elements match, and all comparisons are valid (result in `false`).
            Some(false) => {
                // False, but match the nullability to the input list array.
                Ok(
                    ConstantArray::new(Scalar::bool(false, nullability), list_array.len())
                        .into_array(),
                )
            }
            // All elements match, and all comparisons are valid (result in `true`).
            Some(true) => {
                // True, unless the list itself is empty or NULL.
                list_is_not_empty(&list_array, nullability)
            }
        };
    }

    let ends = list_array.offsets().to_primitive()?;
    match_each_integer_ptype!(ends.ptype(), |T| {
        Ok(reduce_with_ends(
            ends.as_slice::<T>(),
            matches.boolean_buffer(),
            list_array.validity().clone().union_nullability(nullability),
        ))
    })
}

/// Returns a `Bool` array with `false` for lists that are valid,
/// or `NULL` if the list itself is null.
fn list_false_or_null(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {
    match list_array.validity() {
        Validity::NonNullable => {
            // All false.
            Ok(ConstantArray::new(Scalar::bool(false, nullability), list_array.len()).into_array())
        }
        Validity::AllValid => {
            // All false, but nullable.
            Ok(
                ConstantArray::new(Scalar::bool(false, Nullability::Nullable), list_array.len())
                    .into_array(),
            )
        }
        Validity::AllInvalid => {
            // All nulls, must be nullable result.
            Ok(ConstantArray::new(
                Scalar::null(DType::Bool(Nullability::Nullable)),
                list_array.len(),
            )
            .into_array())
        }
        Validity::Array(validity_array) => {
            // Create a new bool array with false, and the provided nulls
            let buffer = BooleanBuffer::new_unset(list_array.len());
            Ok(BoolArray::new(buffer, Validity::Array(validity_array.clone())).into_array())
        }
    }
}

/// Returns a `Bool` array with `true` for lists which are NOT empty, or `false` if they are empty,
/// or `NULL` if the list itself is null.
fn list_is_not_empty(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {
    // Short-circuit for all invalid.
    if matches!(list_array.validity(), Validity::AllInvalid) {
        return Ok(ConstantArray::new(
            Scalar::null(DType::Bool(Nullability::Nullable)),
            list_array.len(),
        )
        .into_array());
    }

    let offsets = list_array.offsets().to_primitive()?;
    let buffer = match_each_integer_ptype!(offsets.ptype(), |T| {
        element_is_not_empty(offsets.as_slice::<T>())
    });

    // Copy over the validity mask from the input.
    Ok(BoolArray::new(
        buffer,
        list_array.validity().clone().union_nullability(nullability),
    )
    .into_array())
}

/// Reduces each boolean values into a Mask that indicates which elements in the
/// ListArray contain the matching value.
fn reduce_with_ends<T: NativePType + AsPrimitive<usize>>(
    ends: &[T],
    matches: &BooleanBuffer,
    validity: Validity,
) -> ArrayRef {
    let mask: BooleanBuffer = ends
        .windows(2)
        .map(|window| {
            let len = window[1].as_() - window[0].as_();
            let mut set_bits = BitIndexIterator::new(matches.values(), window[0].as_(), len);
            set_bits.next().is_some()
        })
        .collect();

    BoolArray::new(mask, validity).into_array()
}

/// Returns a new array of `u64` representing the length of each list element.
///
/// ## Example
///
/// ```rust
/// use vortex_array::arrays::{ListArray, VarBinArray};
/// use vortex_array::{Array, IntoArray};
/// use vortex_array::compute::{list_elem_len};
/// use vortex_array::validity::Validity;
/// use vortex_buffer::buffer;
/// use vortex_dtype::DType;
///
/// let elements = VarBinArray::from_vec(
///         vec!["a", "a", "b", "a", "c"], DType::Utf8(false.into())).into_array();
/// let offsets = buffer![0u32, 1, 3, 5].into_array();
/// let list_array = ListArray::try_new(elements, offsets, Validity::NonNullable).unwrap();
///
/// let lens = list_elem_len(list_array.as_ref()).unwrap();
/// assert_eq!(lens.scalar_at(0).unwrap(), 1u32.into());
/// assert_eq!(lens.scalar_at(1).unwrap(), 2u32.into());
/// assert_eq!(lens.scalar_at(2).unwrap(), 2u32.into());
/// ```
pub fn list_elem_len(array: &dyn Array) -> VortexResult<ArrayRef> {
    if !matches!(array.dtype(), DType::List(..)) {
        vortex_bail!("Array must be of list type");
    }

    // Short-circuit for constant list arrays.
    if array.is_constant() && array.len() > 1 {
        let elem_lens = list_elem_len(&array.slice(0, 1)?)?;
        return Ok(ConstantArray::new(elem_lens.scalar_at(0)?, array.len()).into_array());
    }

    let list_array = array.to_list()?;
    let offsets = list_array.offsets().to_primitive()?;
    let lens_array = match_each_integer_ptype!(offsets.ptype(), |T| {
        element_lens(offsets.as_slice::<T>()).into_array()
    });

    Ok(lens_array)
}

fn element_lens<T: NativePType>(values: &[T]) -> Buffer<T> {
    values
        .windows(2)
        .map(|window| window[1] - window[0])
        .collect()
}

fn element_is_not_empty<T: NativePType>(values: &[T]) -> BooleanBuffer {
    BooleanBuffer::from_iter(values.windows(2).map(|window| window[1] != window[0]))
}

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

    use itertools::Itertools;
    use rstest::rstest;
    use vortex_buffer::Buffer;
    use vortex_dtype::{DType, Nullability, PType};
    use vortex_scalar::Scalar;

    use crate::arrays::{
        BoolArray, ConstantArray, ConstantVTable, ListArray, PrimitiveArray, VarBinArray,
    };
    use crate::canonical::ToCanonical;
    use crate::compute::list_contains;
    use crate::validity::Validity;
    use crate::vtable::ValidityHelper;
    use crate::{Array, ArrayRef, IntoArray};

    fn nonnull_strings(values: Vec<Vec<&str>>) -> ArrayRef {
        ListArray::from_iter_slow::<u64, _>(values, Arc::new(DType::Utf8(Nullability::NonNullable)))
            .unwrap()
    }

    fn null_strings(values: Vec<Vec<Option<&str>>>) -> ArrayRef {
        let elements = values.iter().flatten().cloned().collect_vec();
        let mut offsets = values
            .iter()
            .scan(0u64, |st, v| {
                *st += v.len() as u64;
                Some(*st)
            })
            .collect_vec();
        offsets.insert(0, 0u64);
        let offsets = Buffer::from_iter(offsets).into_array();

        let elements =
            VarBinArray::from_iter(elements, DType::Utf8(Nullability::Nullable)).into_array();

        ListArray::try_new(elements, offsets, Validity::NonNullable)
            .unwrap()
            .into_array()
    }

    fn bool_array(values: Vec<bool>, validity: Validity) -> BoolArray {
        BoolArray::new(values.into_iter().collect(), validity)
    }

    #[rstest]
    #[case(
        nonnull_strings(vec![vec![], vec!["a"], vec!["a", "b"]]),
        Some("a"),
        bool_array(vec![false, true, true], Validity::NonNullable)
    )]
    // Cast 2: valid scalar search over nullable list, with all nulls matched
    #[case(
        null_strings(vec![vec![], vec![Some("a"), None], vec![Some("a"), None, Some("b")]]),
        Some("a"),
        bool_array(vec![false, true, true], Validity::AllValid)
    )]
    // Cast 3: valid scalar search over nullable list, with some nulls not matched (return no nulls)
    #[case(
        null_strings(vec![vec![], vec![Some("a"), None], vec![Some("b"), None, None]]),
        Some("a"),
        bool_array(vec![false, true, false], Validity::AllValid)
    )]
    // Case 4: list(utf8) with all elements matching, but some empty lists
    #[case(
        nonnull_strings(vec![vec![], vec!["a"], vec!["a"]]),
        Some("a"),
        bool_array(vec![false, true, true], Validity::NonNullable)
    )]
    // Case 5: list(utf8) all lists empty.
    #[case(
        nonnull_strings(vec![vec![], vec![], vec![]]),
        Some("a"),
        bool_array(vec![false, false, false], Validity::NonNullable)
    )]
    // Case 6: list(utf8) no elements matching.
    #[case(
        nonnull_strings(vec![vec!["b"], vec![], vec!["b"]]),
        Some("a"),
        bool_array(vec![false, false, false], Validity::NonNullable)
    )]
    // Case 7: list(utf8?) with empty + NULL elements and NULL search
    #[case(
        null_strings(vec![vec![], vec![None, None], vec![None, None, None]]),
        None,
        bool_array(vec![false, true, true], Validity::AllInvalid)
    )]
    // Case 8: list(utf8?) with empty + NULL elements and search scalar
    #[case(
        null_strings(vec![vec![], vec![None, None], vec![None, None, None]]),
        Some("a"),
        bool_array(vec![false, false, false], Validity::AllValid)
    )]
    fn test_contains_nullable(
        #[case] list_array: ArrayRef,
        #[case] value: Option<&str>,
        #[case] expected: BoolArray,
    ) {
        let element_nullability = list_array.dtype().as_list_element().unwrap().nullability();
        let scalar = match value {
            None => Scalar::null(DType::Utf8(Nullability::Nullable)),
            Some(v) => Scalar::utf8(v, element_nullability),
        };
        let elem = ConstantArray::new(scalar, list_array.len());
        let result = list_contains(&list_array, elem.as_ref()).expect("list_contains failed");
        let bool_result = result.to_bool().expect("to_bool failed");
        assert_eq!(
            bool_result.opt_bool_vec().unwrap(),
            expected.opt_bool_vec().unwrap()
        );
        assert_eq!(bool_result.validity(), expected.validity());
    }

    #[test]
    fn test_constant_list() {
        let list_array = ConstantArray::new(
            Scalar::list(
                Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
                vec![1i32.into(), 2i32.into(), 3i32.into()],
                Nullability::NonNullable,
            ),
            2,
        )
        .into_array();

        let contains = list_contains(
            &list_array,
            ConstantArray::new(Scalar::from(2i32), list_array.len()).as_ref(),
        )
        .unwrap();
        assert!(contains.is::<ConstantVTable>(), "Expected constant result");
        assert_eq!(
            contains
                .to_bool()
                .unwrap()
                .boolean_buffer()
                .iter()
                .collect_vec(),
            vec![true, true],
        );
    }

    #[test]
    fn test_all_nulls() {
        let list_array = ConstantArray::new(
            Scalar::null(DType::List(
                Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
                Nullability::Nullable,
            )),
            5,
        )
        .into_array();

        let contains = list_contains(
            &list_array,
            ConstantArray::new(Scalar::from(2i32), list_array.len()).as_ref(),
        )
        .unwrap();
        assert!(contains.is::<ConstantVTable>(), "Expected constant result");

        assert_eq!(contains.len(), 5);
        assert_eq!(
            contains.to_bool().unwrap().validity(),
            &Validity::AllInvalid
        );
    }

    #[test]
    fn test_list_array_element() {
        let list_scalar = Scalar::list(
            Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
            vec![1.into(), 3.into(), 6.into()],
            Nullability::NonNullable,
        );

        let contains = list_contains(
            ConstantArray::new(list_scalar, 7).as_ref(),
            (0..7).collect::<PrimitiveArray>().as_ref(),
        )
        .unwrap();

        assert_eq!(contains.len(), 7);
        assert_eq!(
            contains.to_bool().unwrap().opt_bool_vec().unwrap(),
            vec![
                Some(false),
                Some(true),
                Some(false),
                Some(true),
                Some(false),
                Some(false),
                Some(true)
            ]
        );
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	//! List-related compute operations.
5
6	use std::sync::LazyLock;
7
8	use arcref::ArcRef;
9	use arrow_buffer::BooleanBuffer;
10	use arrow_buffer::bit_iterator::BitIndexIterator;
11	use num_traits::AsPrimitive;
12	use vortex_buffer::Buffer;
13	use vortex_dtype::{DType, NativePType, Nullability, match_each_integer_ptype};
14	use vortex_error::{VortexExpect, VortexResult, vortex_bail};
15	use vortex_scalar::{ListScalar, Scalar};
16
17	use crate::arrays::{BoolArray, ConstantArray, ListArray};
18	use crate::compute::{
19	BinaryArgs, ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Operator, Output, compare,
20	fill_null, or,
21	};
22	use crate::validity::Validity;
23	use crate::vtable::{VTable, ValidityHelper};
24	use crate::{Array, ArrayRef, IntoArray, ToCanonical};
25
26	/// Compute a `Bool`-typed array the same length as `array` where elements is `true` if the list
27	/// item contains the `value`, `false` otherwise.
28	///
29	/// ## Null scalar handling
30	///
31	/// If the `value` or `array` is `null` at any index the result at that index is `null`.
32	///
33	/// ## Format semantics
34	/// ```txt
35	/// list_contains(list, elem)
36	/// ==> (!is_null(list) or NULL) and (!is_null(elem) or NULL) and any({elem = elem_i \| elem_i in list}),
37	/// ```
38	///
39	/// ## Example
40	///
41	/// ```rust
42	/// use vortex_array::{Array, IntoArray, ToCanonical};
43	/// use vortex_array::arrays::{ConstantArray, ListArray, VarBinArray};
44	/// use vortex_array::compute::list_contains;
45	/// use vortex_array::validity::Validity;
46	/// use vortex_buffer::buffer;
47	/// use vortex_dtype::DType;
48	/// use vortex_scalar::Scalar;
49	/// let elements = VarBinArray::from_vec(
50	/// vec!["a", "a", "b", "a", "c"], DType::Utf8(false.into())).into_array();
51	/// let offsets = buffer![0u32, 1, 3, 5].into_array();
52	/// let list_array = ListArray::try_new(elements, offsets, Validity::NonNullable).unwrap();
53	///
54	/// let matches = list_contains(list_array.as_ref(), ConstantArray::new(Scalar::from("b"), list_array.len()).as_ref()).unwrap();
55	/// let to_vec: Vec<bool> = matches.to_bool().unwrap().boolean_buffer().iter().collect();
56	/// assert_eq!(to_vec, vec![false, true, false]);
57	/// ```
58	pub fn list_contains(array: &dyn Array, value: &dyn Array) -> VortexResult<ArrayRef> {	305✔
59	LIST_CONTAINS_FN	305✔
60	.invoke(&InvocationArgs {	305✔
61	inputs: &[array.into(), value.into()],	305✔
62	options: &(),	305✔
63	})?	305✔
64	.unwrap_array()	305✔
65	}	305✔
66
67	pub struct ListContains;
68
69	impl ComputeFnVTable for ListContains {
70	fn invoke(	305✔
71	&self,	305✔
72	args: &InvocationArgs,	305✔
73	kernels: &[ArcRef<dyn Kernel>],	305✔
74	) -> VortexResult<Output> {	305✔
75	let BinaryArgs {
76	lhs: array,	305✔
77	rhs: value,	305✔
78	..
79	} = BinaryArgs::<()>::try_from(args)?;	305✔
80
81	let DType::List(elem_dtype, _) = array.dtype() else {	305✔
82	vortex_bail!("Array must be of List type");	×
83	};
84	if !elem_dtype.as_ref().eq_ignore_nullability(value.dtype()) {	305✔
85	vortex_bail!(	×
86	"Element type {} of ListArray does not match search value {}",	×
87	elem_dtype,
88	value.dtype(),	×
89	);
90	};	305✔
91
92	if value.all_invalid()? \|\| array.all_invalid()? {	305✔
93	return Ok(Output::Array(	2✔
94	ConstantArray::new(	2✔
95	Scalar::null(DType::Bool(Nullability::Nullable)),	2✔
96	array.len(),	2✔
97	)	2✔
98	.to_array(),	2✔
99	));	2✔
100	}	303✔
101
102	for kernel in kernels {	417✔
103	if let Some(output) = kernel.invoke(args)? {	186✔
104	return Ok(output);	72✔
105	}	114✔
106	}
107	if let Some(output) = array.invoke(&LIST_CONTAINS_FN, args)? {	231✔
108	return Ok(output);	×
109	}	231✔
110
111	let nullability = array.dtype().nullability() \| value.dtype().nullability();	231✔
112
113	let result = if let Some(value_scalar) = value.as_constant() {	231✔
114	list_contains_scalar(array, &value_scalar, nullability)	188✔
115	} else if let Some(list_scalar) = array.as_constant() {	43✔
116	constant_list_scalar_contains(&list_scalar.as_list(), value, nullability)	43✔
117	} else {
118	todo!("unsupported list contains with list and element as arrays")	×
119	};
120
121	result.map(Output::Array)	231✔
122	}	305✔
123
124	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	305✔
125	let input = BinaryArgs::<()>::try_from(args)?;	305✔
126	Ok(DType::Bool(	305✔
127	input.lhs.dtype().nullability() \| input.rhs.dtype().nullability(),	305✔
128	))	305✔
129	}	305✔
130
131	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	305✔
132	Ok(BinaryArgs::<()>::try_from(args)?.lhs.len())	305✔
133	}	305✔
134
135	fn is_elementwise(&self) -> bool {	305✔
136	true	305✔
137	}	305✔
138	}
139
140	pub trait ListContainsKernel: VTable {
141	fn list_contains(
142	&self,
143	list: &dyn Array,
144	element: &Self::Array,
145	) -> VortexResult<Option<ArrayRef>>;
146	}
147
148	pub struct ListContainsKernelRef(ArcRef<dyn Kernel>);
149	inventory::collect!(ListContainsKernelRef);
150
151	#[derive(Debug)]
152	pub struct ListContainsKernelAdapter<V: VTable>(pub V);
153
154	impl<V: VTable + ListContainsKernel> ListContainsKernelAdapter<V> {
155	pub const fn lift(&'static self) -> ListContainsKernelRef {	×
156	ListContainsKernelRef(ArcRef::new_ref(self))	×
157	}	×
158	}
159
160	impl<V: VTable + ListContainsKernel> Kernel for ListContainsKernelAdapter<V> {
161	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	74✔
162	let BinaryArgs {
163	lhs: array,	74✔
164	rhs: value,	74✔
165	..
166	} = BinaryArgs::<()>::try_from(args)?;	74✔
167	let Some(value) = value.as_opt::<V>() else {	74✔
168	return Ok(None);	72✔
169	};
170	self.0	2✔
171	.list_contains(array, value)	2✔
172	.map(\|c\| c.map(Output::Array))	2✔
173	}	74✔
174	}
175
176	pub static LIST_CONTAINS_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	265✔
177	let compute = ComputeFn::new("list_contains".into(), ArcRef::new_ref(&ListContains));	265✔
178	for kernel in inventory::iter::<ListContainsKernelRef> {	375✔
179	compute.register_kernel(kernel.0.clone());	110✔
180	}	110✔
181	compute	265✔
182	});	265✔
183
184	// Then there is a constant list scalar (haystack) being compared to an array of needles.
185	fn constant_list_scalar_contains(	43✔
186	list_scalar: &ListScalar,	43✔
187	values: &dyn Array,	43✔
188	nullability: Nullability,	43✔
189	) -> VortexResult<ArrayRef> {	43✔
190	let elements = list_scalar.elements().vortex_expect("non null");	43✔
191
192	let len = values.len();	43✔
193	let mut result: Option<ArrayRef> = None;	43✔
194	let false_scalar = Scalar::bool(false, nullability);	43✔
195	for element in elements {	178✔
196	let res = fill_null(	135✔
197	&compare(	135✔
198	ConstantArray::new(element, len).as_ref(),	135✔
199	values,	135✔
200	Operator::Eq,	135✔
UNCOV 201	)?,	×
202	&false_scalar,	135✔
203	)?;	×
204	if let Some(acc) = result {	135✔
205	result = Some(or(&acc, &res)?)	92✔
206	} else {	43✔
207	result = Some(res);	43✔
208	}	43✔
209	}
210	Ok(result.unwrap_or_else(\|\| ConstantArray::new(false_scalar, len).to_array()))	43✔
211	}	43✔
212
213	fn list_contains_scalar(	190✔
214	array: &dyn Array,	190✔
215	value: &Scalar,	190✔
216	nullability: Nullability,	190✔
217	) -> VortexResult<ArrayRef> {	190✔
218	// If the list array is constant, we perform a single comparison.
219	if array.len() > 1 && array.is_constant() {	190✔
220	let contains = list_contains_scalar(&array.slice(0, 1)?, value, nullability)?;	2✔
221	return Ok(ConstantArray::new(contains.scalar_at(0)?, array.len()).into_array());	2✔
222	}	188✔
223
224	// Canonicalize to a list array.
225	// NOTE(ngates): we may wish to add elements and offsets accessors to the ListArrayTrait.
226	let list_array = array.to_list()?;	188✔
227
228	let elems = list_array.elements();	188✔
229	if elems.is_empty() {	188✔
230	// Must return false when a list is empty (but valid), or null when the list itself is null.
231	return list_false_or_null(&list_array, nullability);	1✔
232	}	187✔
233
234	let rhs = ConstantArray::new(value.clone(), elems.len());	187✔
235	let matching_elements = compare(elems, rhs.as_ref(), Operator::Eq)?;	187✔
236	let matches = matching_elements.to_bool()?;	187✔
237
238	// Fast path: no elements match.
239	if let Some(pred) = matches.as_constant() {	187✔
240	return match pred.as_bool().value() {	39✔
241	// All comparisons are invalid (result in `null`), and search is not null because
242	// we already checked for null above.
243	None => {
244	assert!(	1✔
245	!rhs.scalar().is_null(),	1✔
246	"Search value must not be null here"	×
247	);
248	// False, unless the list itself is null in which case we return null.
249	list_false_or_null(&list_array, nullability)	1✔
250	}
251	// No elements match, and all comparisons are valid (result in `false`).
252	Some(false) => {
253	// False, but match the nullability to the input list array.
254	Ok(	37✔
255	ConstantArray::new(Scalar::bool(false, nullability), list_array.len())	37✔
256	.into_array(),	37✔
257	)	37✔
258	}
259	// All elements match, and all comparisons are valid (result in `true`).
260	Some(true) => {
261	// True, unless the list itself is empty or NULL.
262	list_is_not_empty(&list_array, nullability)	1✔
263	}
264	};
265	}	148✔
266
267	let ends = list_array.offsets().to_primitive()?;	148✔
268	match_each_integer_ptype!(ends.ptype(), \|T\| {	148✔
269	Ok(reduce_with_ends(	×
270	ends.as_slice::<T>(),	×
271	matches.boolean_buffer(),	×
272	list_array.validity().clone().union_nullability(nullability),	×
273	))	×
274	})
275	}	190✔
276
277	/// Returns a `Bool` array with `false` for lists that are valid,
278	/// or `NULL` if the list itself is null.
279	fn list_false_or_null(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {	2✔
280	match list_array.validity() {	2✔
281	Validity::NonNullable => {
282	// All false.
283	Ok(ConstantArray::new(Scalar::bool(false, nullability), list_array.len()).into_array())	2✔
284	}
285	Validity::AllValid => {
286	// All false, but nullable.
287	Ok(	×
288	ConstantArray::new(Scalar::bool(false, Nullability::Nullable), list_array.len())	×
289	.into_array(),	×
290	)	×
291	}
292	Validity::AllInvalid => {
293	// All nulls, must be nullable result.
294	Ok(ConstantArray::new(	×
295	Scalar::null(DType::Bool(Nullability::Nullable)),	×
296	list_array.len(),	×
297	)	×
298	.into_array())	×
299	}
300	Validity::Array(validity_array) => {	×
301	// Create a new bool array with false, and the provided nulls
302	let buffer = BooleanBuffer::new_unset(list_array.len());	×
303	Ok(BoolArray::new(buffer, Validity::Array(validity_array.clone())).into_array())	×
304	}
305	}
306	}	2✔
307
308	/// Returns a `Bool` array with `true` for lists which are NOT empty, or `false` if they are empty,
309	/// or `NULL` if the list itself is null.
310	fn list_is_not_empty(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {	1✔
311	// Short-circuit for all invalid.
312	if matches!(list_array.validity(), Validity::AllInvalid) {	1✔
313	return Ok(ConstantArray::new(	×
314	Scalar::null(DType::Bool(Nullability::Nullable)),	×
315	list_array.len(),	×
316	)	×
317	.into_array());	×
318	}	1✔
319
320	let offsets = list_array.offsets().to_primitive()?;	1✔
321	let buffer = match_each_integer_ptype!(offsets.ptype(), \|T\| {	1✔
322	element_is_not_empty(offsets.as_slice::<T>())	×
323	});
324
325	// Copy over the validity mask from the input.
326	Ok(BoolArray::new(	1✔
327	buffer,	1✔
328	list_array.validity().clone().union_nullability(nullability),	1✔
329	)	1✔
330	.into_array())	1✔
331	}	1✔
332
333	/// Reduces each boolean values into a Mask that indicates which elements in the
334	/// ListArray contain the matching value.
335	fn reduce_with_ends<T: NativePType + AsPrimitive<usize>>(	148✔
336	ends: &[T],	148✔
337	matches: &BooleanBuffer,	148✔
338	validity: Validity,	148✔
339	) -> ArrayRef {	148✔
340	let mask: BooleanBuffer = ends	148✔
341	.windows(2)	148✔
342	.map(\|window\| {	298✔
343	let len = window[1].as_() - window[0].as_();	298✔
344	let mut set_bits = BitIndexIterator::new(matches.values(), window[0].as_(), len);	298✔
345	set_bits.next().is_some()	298✔
346	})	298✔
347	.collect();	148✔
348
349	BoolArray::new(mask, validity).into_array()	148✔
350	}	148✔
351
352	/// Returns a new array of `u64` representing the length of each list element.
353	///
354	/// ## Example
355	///
356	/// ```rust
357	/// use vortex_array::arrays::{ListArray, VarBinArray};
358	/// use vortex_array::{Array, IntoArray};
359	/// use vortex_array::compute::{list_elem_len};
360	/// use vortex_array::validity::Validity;
361	/// use vortex_buffer::buffer;
362	/// use vortex_dtype::DType;
363	///
364	/// let elements = VarBinArray::from_vec(
365	/// vec!["a", "a", "b", "a", "c"], DType::Utf8(false.into())).into_array();
366	/// let offsets = buffer![0u32, 1, 3, 5].into_array();
367	/// let list_array = ListArray::try_new(elements, offsets, Validity::NonNullable).unwrap();
368	///
369	/// let lens = list_elem_len(list_array.as_ref()).unwrap();
370	/// assert_eq!(lens.scalar_at(0).unwrap(), 1u32.into());
371	/// assert_eq!(lens.scalar_at(1).unwrap(), 2u32.into());
372	/// assert_eq!(lens.scalar_at(2).unwrap(), 2u32.into());
373	/// ```
374	pub fn list_elem_len(array: &dyn Array) -> VortexResult<ArrayRef> {	×
375	if !matches!(array.dtype(), DType::List(..)) {	×
376	vortex_bail!("Array must be of list type");	×
377	}	×
378
379	// Short-circuit for constant list arrays.
380	if array.is_constant() && array.len() > 1 {	×
381	let elem_lens = list_elem_len(&array.slice(0, 1)?)?;	×
382	return Ok(ConstantArray::new(elem_lens.scalar_at(0)?, array.len()).into_array());	×
383	}	×
384
385	let list_array = array.to_list()?;	×
386	let offsets = list_array.offsets().to_primitive()?;	×
387	let lens_array = match_each_integer_ptype!(offsets.ptype(), \|T\| {	×
388	element_lens(offsets.as_slice::<T>()).into_array()	×
389	});
390
391	Ok(lens_array)	×
392	}	×
393
394	fn element_lens<T: NativePType>(values: &[T]) -> Buffer<T> {	×
395	values	×
396	.windows(2)	×
397	.map(\|window\| window[1] - window[0])	×
398	.collect()	×
399	}	×
400
401	fn element_is_not_empty<T: NativePType>(values: &[T]) -> BooleanBuffer {	1✔
402	BooleanBuffer::from_iter(values.windows(2).map(\|window\| window[1] != window[0]))	3✔
403	}	1✔
404
405	#[cfg(test)]
406	mod tests {
407	use std::sync::Arc;
408
409	use itertools::Itertools;
410	use rstest::rstest;
411	use vortex_buffer::Buffer;
412	use vortex_dtype::{DType, Nullability, PType};
413	use vortex_scalar::Scalar;
414
415	use crate::arrays::{
416	BoolArray, ConstantArray, ConstantVTable, ListArray, PrimitiveArray, VarBinArray,
417	};
418	use crate::canonical::ToCanonical;
419	use crate::compute::list_contains;
420	use crate::validity::Validity;
421	use crate::vtable::ValidityHelper;
422	use crate::{Array, ArrayRef, IntoArray};
423
424	fn nonnull_strings(values: Vec<Vec<&str>>) -> ArrayRef {	4✔
425	ListArray::from_iter_slow::<u64, _>(values, Arc::new(DType::Utf8(Nullability::NonNullable)))	4✔
426	.unwrap()	4✔
427	}	4✔
428
429	fn null_strings(values: Vec<Vec<Option<&str>>>) -> ArrayRef {	4✔
430	let elements = values.iter().flatten().cloned().collect_vec();	4✔
431	let mut offsets = values	4✔
432	.iter()	4✔
433	.scan(0u64, \|st, v\| {	12✔
434	*st += v.len() as u64;	12✔
435	Some(*st)	12✔
436	})	12✔
437	.collect_vec();	4✔
438	offsets.insert(0, 0u64);	4✔
439	let offsets = Buffer::from_iter(offsets).into_array();	4✔
440
441	let elements =	4✔
442	VarBinArray::from_iter(elements, DType::Utf8(Nullability::Nullable)).into_array();	4✔
443
444	ListArray::try_new(elements, offsets, Validity::NonNullable)	4✔
445	.unwrap()	4✔
446	.into_array()	4✔
447	}	4✔
448
449	fn bool_array(values: Vec<bool>, validity: Validity) -> BoolArray {	8✔
450	BoolArray::new(values.into_iter().collect(), validity)	8✔
451	}	8✔
452
453	#[rstest]
454	#[case(
455	nonnull_strings(vec![vec![], vec!["a"], vec!["a", "b"]]),
456	Some("a"),
457	bool_array(vec![false, true, true], Validity::NonNullable)
458	)]
459	// Cast 2: valid scalar search over nullable list, with all nulls matched
460	#[case(
461	null_strings(vec![vec![], vec![Some("a"), None], vec![Some("a"), None, Some("b")]]),
462	Some("a"),
463	bool_array(vec![false, true, true], Validity::AllValid)
464	)]
465	// Cast 3: valid scalar search over nullable list, with some nulls not matched (return no nulls)
466	#[case(
467	null_strings(vec![vec![], vec![Some("a"), None], vec![Some("b"), None, None]]),
468	Some("a"),
469	bool_array(vec![false, true, false], Validity::AllValid)
470	)]
471	// Case 4: list(utf8) with all elements matching, but some empty lists
472	#[case(
473	nonnull_strings(vec![vec![], vec!["a"], vec!["a"]]),
474	Some("a"),
475	bool_array(vec![false, true, true], Validity::NonNullable)
476	)]
477	// Case 5: list(utf8) all lists empty.
478	#[case(
479	nonnull_strings(vec![vec![], vec![], vec![]]),
480	Some("a"),
481	bool_array(vec![false, false, false], Validity::NonNullable)
482	)]
483	// Case 6: list(utf8) no elements matching.
484	#[case(
485	nonnull_strings(vec![vec!["b"], vec![], vec!["b"]]),
486	Some("a"),
487	bool_array(vec![false, false, false], Validity::NonNullable)
488	)]
489	// Case 7: list(utf8?) with empty + NULL elements and NULL search
490	#[case(
491	null_strings(vec![vec![], vec![None, None], vec![None, None, None]]),
492	None,
493	bool_array(vec![false, true, true], Validity::AllInvalid)
494	)]
495	// Case 8: list(utf8?) with empty + NULL elements and search scalar
496	#[case(
497	null_strings(vec![vec![], vec![None, None], vec![None, None, None]]),
498	Some("a"),
499	bool_array(vec![false, false, false], Validity::AllValid)
500	)]
501	fn test_contains_nullable(
502	#[case] list_array: ArrayRef,
503	#[case] value: Option<&str>,
504	#[case] expected: BoolArray,
505	) {
506	let element_nullability = list_array.dtype().as_list_element().unwrap().nullability();
507	let scalar = match value {
508	None => Scalar::null(DType::Utf8(Nullability::Nullable)),
509	Some(v) => Scalar::utf8(v, element_nullability),
510	};
511	let elem = ConstantArray::new(scalar, list_array.len());
512	let result = list_contains(&list_array, elem.as_ref()).expect("list_contains failed");
513	let bool_result = result.to_bool().expect("to_bool failed");
514	assert_eq!(
515	bool_result.opt_bool_vec().unwrap(),
516	expected.opt_bool_vec().unwrap()
517	);
518	assert_eq!(bool_result.validity(), expected.validity());
519	}
520
521	#[test]
522	fn test_constant_list() {	1✔
523	let list_array = ConstantArray::new(	1✔
524	Scalar::list(	1✔
525	Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),	1✔
526	vec![1i32.into(), 2i32.into(), 3i32.into()],	1✔
527	Nullability::NonNullable,	1✔
528	),
529	2,
530	)
531	.into_array();	1✔
532
533	let contains = list_contains(	1✔
534	&list_array,	1✔
535	ConstantArray::new(Scalar::from(2i32), list_array.len()).as_ref(),	1✔
536	)
537	.unwrap();	1✔
538	assert!(contains.is::<ConstantVTable>(), "Expected constant result");	1✔
539	assert_eq!(	1✔
540	contains	1✔
541	.to_bool()	1✔
542	.unwrap()	1✔
543	.boolean_buffer()	1✔
544	.iter()	1✔
545	.collect_vec(),	1✔
546	vec![true, true],	1✔
547	);
548	}	1✔
549
550	#[test]
551	fn test_all_nulls() {	1✔
552	let list_array = ConstantArray::new(	1✔
553	Scalar::null(DType::List(	1✔
554	Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),	1✔
555	Nullability::Nullable,	1✔
556	)),	1✔
557	5,
558	)
559	.into_array();	1✔
560
561	let contains = list_contains(	1✔
562	&list_array,	1✔
563	ConstantArray::new(Scalar::from(2i32), list_array.len()).as_ref(),	1✔
564	)
565	.unwrap();	1✔
566	assert!(contains.is::<ConstantVTable>(), "Expected constant result");	1✔
567
568	assert_eq!(contains.len(), 5);	1✔
569	assert_eq!(	1✔
570	contains.to_bool().unwrap().validity(),	1✔
571	&Validity::AllInvalid
572	);
573	}	1✔
574
575	#[test]
576	fn test_list_array_element() {	1✔
577	let list_scalar = Scalar::list(	1✔
578	Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),	1✔
579	vec![1.into(), 3.into(), 6.into()],	1✔
580	Nullability::NonNullable,	1✔
581	);
582
583	let contains = list_contains(	1✔
584	ConstantArray::new(list_scalar, 7).as_ref(),	1✔
585	(0..7).collect::<PrimitiveArray>().as_ref(),	1✔
586	)
587	.unwrap();	1✔
588
589	assert_eq!(contains.len(), 7);	1✔
590	assert_eq!(	1✔
591	contains.to_bool().unwrap().opt_bool_vec().unwrap(),	1✔
592	vec![	1✔
593	Some(false),	1✔
594	Some(true),	1✔
595	Some(false),	1✔
596	Some(true),	1✔
597	Some(false),	1✔
598	Some(false),	1✔
599	Some(true)	1✔
600	]
601	);
602	}	1✔
603	}

vortex-data / vortex / 16331938722

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous