16992684502

Committed 15 Aug 2025 02:56PM UTC coverage: 87.875% (+0.2%) from 87.72%

Build # 16992684502

Build Type

Pull #2456

github

Committed by

web-flow

Commit Message

Merge 2d540e578 into 4a23f65b3

Pull Request Pull Request #2456: feat: basic BoolBuffer / BoolBufferMut

Run Details

1275 of 1428 new or added lines in 110 files covered. (89.29%)

334 existing lines in 31 files now uncovered.

57169 of 65057 relevant lines covered (87.88%)

658056.52 hits per line

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

81.36

/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::bit_iterator::BitIndexIterator;
use num_traits::AsPrimitive;
use vortex_buffer::{BitBuffer, 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};

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
});

/// 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().bit_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))
    }
}

// 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.bit_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 = BitBuffer::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: &BitBuffer,
    validity: Validity,
) -> ArrayRef {
    let mask: BitBuffer = ends
        .windows(2)
        .map(|window| {
            let len = window[1].as_() - window[0].as_();
            let mut set_bits =
                BitIndexIterator::new(matches.inner().as_slice(), 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]) -> BitBuffer {
    BitBuffer::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()
                .bit_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::bit_iterator::BitIndexIterator;
10	use num_traits::AsPrimitive;
11	use vortex_buffer::{BitBuffer, Buffer};
12	use vortex_dtype::{DType, NativePType, Nullability, match_each_integer_ptype};
13	use vortex_error::{VortexExpect, VortexResult, vortex_bail};
14	use vortex_scalar::{ListScalar, Scalar};
15
16	use crate::arrays::{BoolArray, ConstantArray, ListArray};
17	use crate::compute::{
18	BinaryArgs, ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Operator, Output, compare,
19	fill_null, or,
20	};
21	use crate::validity::Validity;
22	use crate::vtable::{VTable, ValidityHelper};
23	use crate::{Array, ArrayRef, IntoArray, ToCanonical};
24
25	static LIST_CONTAINS_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	287✔
26	let compute = ComputeFn::new("list_contains".into(), ArcRef::new_ref(&ListContains));	287✔
27	for kernel in inventory::iter::<ListContainsKernelRef> {	407✔
28	compute.register_kernel(kernel.0.clone());	120✔
29	}	120✔
30	compute	287✔
31	});	287✔
32
33	/// Compute a `Bool`-typed array the same length as `array` where elements is `true` if the list
34	/// item contains the `value`, `false` otherwise.
35	///
36	/// ## Null scalar handling
37	///
38	/// If the `value` or `array` is `null` at any index the result at that index is `null`.
39	///
40	/// ## Format semantics
41	/// ```txt
42	/// list_contains(list, elem)
43	/// ==> (!is_null(list) or NULL) and (!is_null(elem) or NULL) and any({elem = elem_i \| elem_i in list}),
44	/// ```
45	///
46	/// ## Example
47	///
48	/// ```rust
49	/// use vortex_array::{Array, IntoArray, ToCanonical};
50	/// use vortex_array::arrays::{ConstantArray, ListArray, VarBinArray};
51	/// use vortex_array::compute::list_contains;
52	/// use vortex_array::validity::Validity;
53	/// use vortex_buffer::buffer;
54	/// use vortex_dtype::DType;
55	/// use vortex_scalar::Scalar;
56	/// let elements = VarBinArray::from_vec(
57	/// vec!["a", "a", "b", "a", "c"], DType::Utf8(false.into())).into_array();
58	/// let offsets = buffer![0u32, 1, 3, 5].into_array();
59	/// let list_array = ListArray::try_new(elements, offsets, Validity::NonNullable).unwrap();
60	///
61	/// let matches = list_contains(list_array.as_ref(), ConstantArray::new(Scalar::from("b"), list_array.len()).as_ref()).unwrap();
62	/// let to_vec: Vec<bool> = matches.to_bool().unwrap().bit_buffer().iter().collect();
63	/// assert_eq!(to_vec, vec![false, true, false]);
64	/// ```
65	pub fn list_contains(array: &dyn Array, value: &dyn Array) -> VortexResult<ArrayRef> {	410✔
66	LIST_CONTAINS_FN	410✔
67	.invoke(&InvocationArgs {	410✔
68	inputs: &[array.into(), value.into()],	410✔
69	options: &(),	410✔
70	})?	410✔
71	.unwrap_array()	410✔
72	}	410✔
73
74	pub struct ListContains;
75
76	impl ComputeFnVTable for ListContains {
77	fn invoke(	410✔
78	&self,	410✔
79	args: &InvocationArgs,	410✔
80	kernels: &[ArcRef<dyn Kernel>],	410✔
81	) -> VortexResult<Output> {	410✔
82	let BinaryArgs {
83	lhs: array,	410✔
84	rhs: value,	410✔
85	..
86	} = BinaryArgs::<()>::try_from(args)?;	410✔
87
88	let DType::List(elem_dtype, _) = array.dtype() else {	410✔
89	vortex_bail!("Array must be of List type");	×
90	};
91	if !elem_dtype.as_ref().eq_ignore_nullability(value.dtype()) {	410✔
92	vortex_bail!(	×
93	"Element type {} of ListArray does not match search value {}",	×
94	elem_dtype,
95	value.dtype(),	×
96	);
97	};	410✔
98
99	if value.all_invalid()? \|\| array.all_invalid()? {	410✔
100	return Ok(Output::Array(	2✔
101	ConstantArray::new(	2✔
102	Scalar::null(DType::Bool(Nullability::Nullable)),	2✔
103	array.len(),	2✔
104	)	2✔
105	.to_array(),	2✔
106	));	2✔
107	}	408✔
108
109	for kernel in kernels {	612✔
110	if let Some(output) = kernel.invoke(args)? {	282✔
111	return Ok(output);	78✔
112	}	204✔
113	}
114	if let Some(output) = array.invoke(&LIST_CONTAINS_FN, args)? {	330✔
115	return Ok(output);	×
116	}	330✔
117
118	let nullability = array.dtype().nullability() \| value.dtype().nullability();	330✔
119
120	let result = if let Some(value_scalar) = value.as_constant() {	330✔
121	list_contains_scalar(array, &value_scalar, nullability)	203✔
122	} else if let Some(list_scalar) = array.as_constant() {	127✔
123	constant_list_scalar_contains(&list_scalar.as_list(), value, nullability)	127✔
124	} else {
125	todo!("unsupported list contains with list and element as arrays")	×
126	};
127
128	result.map(Output::Array)	330✔
129	}	410✔
130
131	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	410✔
132	let input = BinaryArgs::<()>::try_from(args)?;	410✔
133	Ok(DType::Bool(	410✔
134	input.lhs.dtype().nullability() \| input.rhs.dtype().nullability(),	410✔
135	))	410✔
136	}	410✔
137
138	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	410✔
139	Ok(BinaryArgs::<()>::try_from(args)?.lhs.len())	410✔
140	}	410✔
141
142	fn is_elementwise(&self) -> bool {	410✔
143	true	410✔
144	}	410✔
145	}
146
147	pub trait ListContainsKernel: VTable {
148	fn list_contains(
149	&self,
150	list: &dyn Array,
151	element: &Self::Array,
152	) -> VortexResult<Option<ArrayRef>>;
153	}
154
155	pub struct ListContainsKernelRef(ArcRef<dyn Kernel>);
156	inventory::collect!(ListContainsKernelRef);
157
158	#[derive(Debug)]
159	pub struct ListContainsKernelAdapter<V: VTable>(pub V);
160
161	impl<V: VTable + ListContainsKernel> ListContainsKernelAdapter<V> {
162	pub const fn lift(&'static self) -> ListContainsKernelRef {	×
163	ListContainsKernelRef(ArcRef::new_ref(self))	×
164	}	×
165	}
166
167	impl<V: VTable + ListContainsKernel> Kernel for ListContainsKernelAdapter<V> {
168	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	161✔
169	let BinaryArgs {
170	lhs: array,	161✔
171	rhs: value,	161✔
172	..
173	} = BinaryArgs::<()>::try_from(args)?;	161✔
174	let Some(value) = value.as_opt::<V>() else {	161✔
175	return Ok(None);	159✔
176	};
177	self.0	2✔
178	.list_contains(array, value)	2✔
179	.map(\|c\| c.map(Output::Array))	2✔
180	}	161✔
181	}
182
183	// Then there is a constant list scalar (haystack) being compared to an array of needles.
184	fn constant_list_scalar_contains(	127✔
185	list_scalar: &ListScalar,	127✔
186	values: &dyn Array,	127✔
187	nullability: Nullability,	127✔
188	) -> VortexResult<ArrayRef> {	127✔
189	let elements = list_scalar.elements().vortex_expect("non null");	127✔
190
191	let len = values.len();	127✔
192	let mut result: Option<ArrayRef> = None;	127✔
193	let false_scalar = Scalar::bool(false, nullability);	127✔
194	for element in elements {	1,837✔
195	let res = fill_null(	1,710✔
196	&compare(	1,710✔
197	ConstantArray::new(element, len).as_ref(),	1,710✔
198	values,	1,710✔
199	Operator::Eq,	1,710✔
200	)?,	×
201	&false_scalar,	1,710✔
202	)?;	×
203	if let Some(acc) = result {	1,710✔
204	result = Some(or(&acc, &res)?)	1,583✔
205	} else {	127✔
206	result = Some(res);	127✔
207	}	127✔
208	}
209	Ok(result.unwrap_or_else(\|\| ConstantArray::new(false_scalar, len).to_array()))	127✔
210	}	127✔
211
212	fn list_contains_scalar(	205✔
213	array: &dyn Array,	205✔
214	value: &Scalar,	205✔
215	nullability: Nullability,	205✔
216	) -> VortexResult<ArrayRef> {	205✔
217	// If the list array is constant, we perform a single comparison.
218	if array.len() > 1 && array.is_constant() {	205✔
219	let contains = list_contains_scalar(&array.slice(0, 1)?, value, nullability)?;	2✔
220	return Ok(ConstantArray::new(contains.scalar_at(0)?, array.len()).into_array());	2✔
221	}	203✔
222
223	// Canonicalize to a list array.
224	// NOTE(ngates): we may wish to add elements and offsets accessors to the ListArrayTrait.
225	let list_array = array.to_list()?;	203✔
226
227	let elems = list_array.elements();	203✔
228	if elems.is_empty() {	203✔
229	// Must return false when a list is empty (but valid), or null when the list itself is null.
230	return list_false_or_null(&list_array, nullability);	1✔
231	}	202✔
232
233	let rhs = ConstantArray::new(value.clone(), elems.len());	202✔
234	let matching_elements = compare(elems, rhs.as_ref(), Operator::Eq)?;	202✔
235	let matches = matching_elements.to_bool()?;	202✔
236
237	// Fast path: no elements match.
238	if let Some(pred) = matches.as_constant() {	202✔
239	return match pred.as_bool().value() {	42✔
240	// All comparisons are invalid (result in `null`), and search is not null because
241	// we already checked for null above.
242	None => {
243	assert!(	1✔
244	!rhs.scalar().is_null(),	1✔
245	"Search value must not be null here"	×
246	);
247	// False, unless the list itself is null in which case we return null.
248	list_false_or_null(&list_array, nullability)	1✔
249	}
250	// No elements match, and all comparisons are valid (result in `false`).
251	Some(false) => {
252	// False, but match the nullability to the input list array.
253	Ok(	40✔
254	ConstantArray::new(Scalar::bool(false, nullability), list_array.len())	40✔
255	.into_array(),	40✔
256	)	40✔
257	}
258	// All elements match, and all comparisons are valid (result in `true`).
259	Some(true) => {
260	// True, unless the list itself is empty or NULL.
261	list_is_not_empty(&list_array, nullability)	1✔
262	}
263	};
264	}	160✔
265
266	let ends = list_array.offsets().to_primitive()?;	160✔
267	match_each_integer_ptype!(ends.ptype(), \|T\| {	160✔
268	Ok(reduce_with_ends(	×
269	ends.as_slice::<T>(),	×
NEW 270	matches.bit_buffer(),	×
271	list_array.validity().clone().union_nullability(nullability),	×
272	))	×
273	})
274	}	205✔
275
276	/// Returns a `Bool` array with `false` for lists that are valid,
277	/// or `NULL` if the list itself is null.
278	fn list_false_or_null(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {	2✔
279	match list_array.validity() {	2✔
280	Validity::NonNullable => {
281	// All false.
282	Ok(ConstantArray::new(Scalar::bool(false, nullability), list_array.len()).into_array())	2✔
283	}
284	Validity::AllValid => {
285	// All false, but nullable.
286	Ok(	×
287	ConstantArray::new(Scalar::bool(false, Nullability::Nullable), list_array.len())	×
288	.into_array(),	×
289	)	×
290	}
291	Validity::AllInvalid => {
292	// All nulls, must be nullable result.
293	Ok(ConstantArray::new(	×
294	Scalar::null(DType::Bool(Nullability::Nullable)),	×
295	list_array.len(),	×
296	)	×
297	.into_array())	×
298	}
299	Validity::Array(validity_array) => {	×
300	// Create a new bool array with false, and the provided nulls
NEW 301	let buffer = BitBuffer::new_unset(list_array.len());	×
302	Ok(BoolArray::new(buffer, Validity::Array(validity_array.clone())).into_array())	×
303	}
304	}
305	}	2✔
306
307	/// Returns a `Bool` array with `true` for lists which are NOT empty, or `false` if they are empty,
308	/// or `NULL` if the list itself is null.
309	fn list_is_not_empty(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {	1✔
310	// Short-circuit for all invalid.
311	if matches!(list_array.validity(), Validity::AllInvalid) {	1✔
312	return Ok(ConstantArray::new(	×
313	Scalar::null(DType::Bool(Nullability::Nullable)),	×
314	list_array.len(),	×
315	)	×
316	.into_array());	×
317	}	1✔
318
319	let offsets = list_array.offsets().to_primitive()?;	1✔
320	let buffer = match_each_integer_ptype!(offsets.ptype(), \|T\| {	1✔
321	element_is_not_empty(offsets.as_slice::<T>())	×
322	});
323
324	// Copy over the validity mask from the input.
325	Ok(BoolArray::new(	1✔
326	buffer,	1✔
327	list_array.validity().clone().union_nullability(nullability),	1✔
328	)	1✔
329	.into_array())	1✔
330	}	1✔
331
332	/// Reduces each boolean values into a Mask that indicates which elements in the
333	/// ListArray contain the matching value.
334	fn reduce_with_ends<T: NativePType + AsPrimitive<usize>>(	160✔
335	ends: &[T],	160✔
336	matches: &BitBuffer,	160✔
337	validity: Validity,	160✔
338	) -> ArrayRef {	160✔
339	let mask: BitBuffer = ends	160✔
340	.windows(2)	160✔
341	.map(\|window\| {	322✔
342	let len = window[1].as_() - window[0].as_();	322✔
343	let mut set_bits =	322✔
344	BitIndexIterator::new(matches.inner().as_slice(), window[0].as_(), len);	322✔
345	set_bits.next().is_some()	322✔
346	})	322✔
347	.collect();	160✔
348
349	BoolArray::new(mask, validity).into_array()	160✔
350	}	160✔
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]) -> BitBuffer {	1✔
402	BitBuffer::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	.bit_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 / 16992684502

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