16935267080

Committed 13 Aug 2025 11:00AM UTC coverage: 24.312% (-63.3%) from 87.658%

Build # 16935267080

Build Type

Pull #4226

github

Committed by

web-flow

Commit Message

Merge 81b48c7fb into baa6ea202

Pull Request Pull Request #4226: Support converting TimestampTZ to and from duckdb

Run Details

0 of 2 new or added lines in 1 file covered. (0.0%)

20666 existing lines in 469 files now uncovered.

8726 of 35892 relevant lines covered (24.31%)

147.74 hits per line

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

0.0

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

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

// 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
UNCOV 26	static LIST_CONTAINS_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	×
UNCOV 27	let compute = ComputeFn::new("list_contains".into(), ArcRef::new_ref(&ListContains));	×
UNCOV 28	for kernel in inventory::iter::<ListContainsKernelRef> {	×
UNCOV 29	compute.register_kernel(kernel.0.clone());	×
UNCOV 30	}	×
UNCOV 31	compute	×
UNCOV 32	});	×
33
34	/// Compute a `Bool`-typed array the same length as `array` where elements is `true` if the list
35	/// item contains the `value`, `false` otherwise.
36	///
37	/// ## Null scalar handling
38	///
39	/// If the `value` or `array` is `null` at any index the result at that index is `null`.
40	///
41	/// ## Format semantics
42	/// ```txt
43	/// list_contains(list, elem)
44	/// ==> (!is_null(list) or NULL) and (!is_null(elem) or NULL) and any({elem = elem_i \| elem_i in list}),
45	/// ```
46	///
47	/// ## Example
48	///
49	/// ```rust
50	/// use vortex_array::{Array, IntoArray, ToCanonical};
51	/// use vortex_array::arrays::{ConstantArray, ListArray, VarBinArray};
52	/// use vortex_array::compute::list_contains;
53	/// use vortex_array::validity::Validity;
54	/// use vortex_buffer::buffer;
55	/// use vortex_dtype::DType;
56	/// use vortex_scalar::Scalar;
57	/// let elements = VarBinArray::from_vec(
58	/// vec!["a", "a", "b", "a", "c"], DType::Utf8(false.into())).into_array();
59	/// let offsets = buffer![0u32, 1, 3, 5].into_array();
60	/// let list_array = ListArray::try_new(elements, offsets, Validity::NonNullable).unwrap();
61	///
62	/// let matches = list_contains(list_array.as_ref(), ConstantArray::new(Scalar::from("b"), list_array.len()).as_ref()).unwrap();
63	/// let to_vec: Vec<bool> = matches.to_bool().unwrap().boolean_buffer().iter().collect();
64	/// assert_eq!(to_vec, vec![false, true, false]);
65	/// ```
UNCOV 66	pub fn list_contains(array: &dyn Array, value: &dyn Array) -> VortexResult<ArrayRef> {	×
UNCOV 67	LIST_CONTAINS_FN	×
UNCOV 68	.invoke(&InvocationArgs {	×
UNCOV 69	inputs: &[array.into(), value.into()],	×
UNCOV 70	options: &(),	×
UNCOV 71	})?	×
UNCOV 72	.unwrap_array()	×
UNCOV 73	}	×
74
75	pub struct ListContains;
76
77	impl ComputeFnVTable for ListContains {
UNCOV 78	fn invoke(	×
UNCOV 79	&self,	×
UNCOV 80	args: &InvocationArgs,	×
UNCOV 81	kernels: &[ArcRef<dyn Kernel>],	×
UNCOV 82	) -> VortexResult<Output> {	×
83	let BinaryArgs {
UNCOV 84	lhs: array,	×
UNCOV 85	rhs: value,	×
86	..
UNCOV 87	} = BinaryArgs::<()>::try_from(args)?;	×
88
UNCOV 89	let DType::List(elem_dtype, _) = array.dtype() else {	×
90	vortex_bail!("Array must be of List type");	×
91	};
UNCOV 92	if !elem_dtype.as_ref().eq_ignore_nullability(value.dtype()) {	×
93	vortex_bail!(	×
94	"Element type {} of ListArray does not match search value {}",	×
95	elem_dtype,
96	value.dtype(),	×
97	);
UNCOV 98	};	×
99
UNCOV 100	if value.all_invalid()? \|\| array.all_invalid()? {	×
UNCOV 101	return Ok(Output::Array(	×
UNCOV 102	ConstantArray::new(	×
UNCOV 103	Scalar::null(DType::Bool(Nullability::Nullable)),	×
UNCOV 104	array.len(),	×
UNCOV 105	)	×
UNCOV 106	.to_array(),	×
UNCOV 107	));	×
UNCOV 108	}	×
109
UNCOV 110	for kernel in kernels {	×
UNCOV 111	if let Some(output) = kernel.invoke(args)? {	×
UNCOV 112	return Ok(output);	×
UNCOV 113	}	×
114	}
UNCOV 115	if let Some(output) = array.invoke(&LIST_CONTAINS_FN, args)? {	×
116	return Ok(output);	×
UNCOV 117	}	×
118
UNCOV 119	let nullability = array.dtype().nullability() \| value.dtype().nullability();	×
120
UNCOV 121	let result = if let Some(value_scalar) = value.as_constant() {	×
UNCOV 122	list_contains_scalar(array, &value_scalar, nullability)	×
UNCOV 123	} else if let Some(list_scalar) = array.as_constant() {	×
UNCOV 124	constant_list_scalar_contains(&list_scalar.as_list(), value, nullability)	×
125	} else {
126	todo!("unsupported list contains with list and element as arrays")	×
127	};
128
UNCOV 129	result.map(Output::Array)	×
UNCOV 130	}	×
131
UNCOV 132	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	×
UNCOV 133	let input = BinaryArgs::<()>::try_from(args)?;	×
UNCOV 134	Ok(DType::Bool(	×
UNCOV 135	input.lhs.dtype().nullability() \| input.rhs.dtype().nullability(),	×
UNCOV 136	))	×
UNCOV 137	}	×
138
UNCOV 139	fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {	×
UNCOV 140	Ok(BinaryArgs::<()>::try_from(args)?.lhs.len())	×
UNCOV 141	}	×
142
UNCOV 143	fn is_elementwise(&self) -> bool {	×
UNCOV 144	true	×
UNCOV 145	}	×
146	}
147
148	pub trait ListContainsKernel: VTable {
149	fn list_contains(
150	&self,
151	list: &dyn Array,
152	element: &Self::Array,
153	) -> VortexResult<Option<ArrayRef>>;
154	}
155
156	pub struct ListContainsKernelRef(ArcRef<dyn Kernel>);
157	inventory::collect!(ListContainsKernelRef);
158
159	#[derive(Debug)]
160	pub struct ListContainsKernelAdapter<V: VTable>(pub V);
161
162	impl<V: VTable + ListContainsKernel> ListContainsKernelAdapter<V> {
163	pub const fn lift(&'static self) -> ListContainsKernelRef {	×
164	ListContainsKernelRef(ArcRef::new_ref(self))	×
165	}	×
166	}
167
168	impl<V: VTable + ListContainsKernel> Kernel for ListContainsKernelAdapter<V> {
UNCOV 169	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	×
170	let BinaryArgs {
UNCOV 171	lhs: array,	×
UNCOV 172	rhs: value,	×
173	..
UNCOV 174	} = BinaryArgs::<()>::try_from(args)?;	×
UNCOV 175	let Some(value) = value.as_opt::<V>() else {	×
UNCOV 176	return Ok(None);	×
177	};
UNCOV 178	self.0	×
UNCOV 179	.list_contains(array, value)	×
UNCOV 180	.map(\|c\| c.map(Output::Array))	×
UNCOV 181	}	×
182	}
183
184	// Then there is a constant list scalar (haystack) being compared to an array of needles.
UNCOV 185	fn constant_list_scalar_contains(	×
UNCOV 186	list_scalar: &ListScalar,	×
UNCOV 187	values: &dyn Array,	×
UNCOV 188	nullability: Nullability,	×
UNCOV 189	) -> VortexResult<ArrayRef> {	×
UNCOV 190	let elements = list_scalar.elements().vortex_expect("non null");	×
191
UNCOV 192	let len = values.len();	×
UNCOV 193	let mut result: Option<ArrayRef> = None;	×
UNCOV 194	let false_scalar = Scalar::bool(false, nullability);	×
UNCOV 195	for element in elements {	×
UNCOV 196	let res = fill_null(	×
UNCOV 197	&compare(	×
UNCOV 198	ConstantArray::new(element, len).as_ref(),	×
UNCOV 199	values,	×
UNCOV 200	Operator::Eq,	×
201	)?,	×
UNCOV 202	&false_scalar,	×
203	)?;	×
UNCOV 204	if let Some(acc) = result {	×
UNCOV 205	result = Some(or(&acc, &res)?)	×
UNCOV 206	} else {	×
UNCOV 207	result = Some(res);	×
UNCOV 208	}	×
209	}
UNCOV 210	Ok(result.unwrap_or_else(\|\| ConstantArray::new(false_scalar, len).to_array()))	×
UNCOV 211	}	×
212
UNCOV 213	fn list_contains_scalar(	×
UNCOV 214	array: &dyn Array,	×
UNCOV 215	value: &Scalar,	×
UNCOV 216	nullability: Nullability,	×
UNCOV 217	) -> VortexResult<ArrayRef> {	×
218	// If the list array is constant, we perform a single comparison.
UNCOV 219	if array.len() > 1 && array.is_constant() {	×
UNCOV 220	let contains = list_contains_scalar(&array.slice(0, 1)?, value, nullability)?;	×
UNCOV 221	return Ok(ConstantArray::new(contains.scalar_at(0)?, array.len()).into_array());	×
UNCOV 222	}	×
223
224	// Canonicalize to a list array.
225	// NOTE(ngates): we may wish to add elements and offsets accessors to the ListArrayTrait.
UNCOV 226	let list_array = array.to_list()?;	×
227
UNCOV 228	let elems = list_array.elements();	×
UNCOV 229	if elems.is_empty() {	×
230	// Must return false when a list is empty (but valid), or null when the list itself is null.
UNCOV 231	return list_false_or_null(&list_array, nullability);	×
UNCOV 232	}	×
233
UNCOV 234	let rhs = ConstantArray::new(value.clone(), elems.len());	×
UNCOV 235	let matching_elements = compare(elems, rhs.as_ref(), Operator::Eq)?;	×
UNCOV 236	let matches = matching_elements.to_bool()?;	×
237
238	// Fast path: no elements match.
UNCOV 239	if let Some(pred) = matches.as_constant() {	×
UNCOV 240	return match pred.as_bool().value() {	×
241	// All comparisons are invalid (result in `null`), and search is not null because
242	// we already checked for null above.
243	None => {
UNCOV 244	assert!(	×
UNCOV 245	!rhs.scalar().is_null(),	×
246	"Search value must not be null here"	×
247	);
248	// False, unless the list itself is null in which case we return null.
UNCOV 249	list_false_or_null(&list_array, nullability)	×
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.
UNCOV 254	Ok(	×
UNCOV 255	ConstantArray::new(Scalar::bool(false, nullability), list_array.len())	×
UNCOV 256	.into_array(),	×
UNCOV 257	)	×
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.
UNCOV 262	list_is_not_empty(&list_array, nullability)	×
263	}
264	};
UNCOV 265	}	×
266
UNCOV 267	let ends = list_array.offsets().to_primitive()?;	×
UNCOV 268	match_each_integer_ptype!(ends.ptype(), \|T\| {	×
269	Ok(reduce_with_ends(	×
270	ends.as_slice::<T>(),	×
271	matches.boolean_buffer(),	×
272	list_array.validity().clone().union_nullability(nullability),	×
273	))	×
274	})
UNCOV 275	}	×
276
277	/// Returns a `Bool` array with `false` for lists that are valid,
278	/// or `NULL` if the list itself is null.
UNCOV 279	fn list_false_or_null(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {	×
UNCOV 280	match list_array.validity() {	×
281	Validity::NonNullable => {
282	// All false.
UNCOV 283	Ok(ConstantArray::new(Scalar::bool(false, nullability), list_array.len()).into_array())	×
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	}
UNCOV 306	}	×
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.
UNCOV 310	fn list_is_not_empty(list_array: &ListArray, nullability: Nullability) -> VortexResult<ArrayRef> {	×
311	// Short-circuit for all invalid.
UNCOV 312	if matches!(list_array.validity(), Validity::AllInvalid) {	×
313	return Ok(ConstantArray::new(	×
314	Scalar::null(DType::Bool(Nullability::Nullable)),	×
315	list_array.len(),	×
316	)	×
317	.into_array());	×
UNCOV 318	}	×
319
UNCOV 320	let offsets = list_array.offsets().to_primitive()?;	×
UNCOV 321	let buffer = match_each_integer_ptype!(offsets.ptype(), \|T\| {	×
322	element_is_not_empty(offsets.as_slice::<T>())	×
323	});
324
325	// Copy over the validity mask from the input.
UNCOV 326	Ok(BoolArray::new(	×
UNCOV 327	buffer,	×
UNCOV 328	list_array.validity().clone().union_nullability(nullability),	×
UNCOV 329	)	×
UNCOV 330	.into_array())	×
UNCOV 331	}	×
332
333	/// Reduces each boolean values into a Mask that indicates which elements in the
334	/// ListArray contain the matching value.
UNCOV 335	fn reduce_with_ends<T: NativePType + AsPrimitive<usize>>(	×
UNCOV 336	ends: &[T],	×
UNCOV 337	matches: &BooleanBuffer,	×
UNCOV 338	validity: Validity,	×
UNCOV 339	) -> ArrayRef {	×
UNCOV 340	let mask: BooleanBuffer = ends	×
UNCOV 341	.windows(2)	×
UNCOV 342	.map(\|window\| {	×
UNCOV 343	let len = window[1].as_() - window[0].as_();	×
UNCOV 344	let mut set_bits = BitIndexIterator::new(matches.values(), window[0].as_(), len);	×
UNCOV 345	set_bits.next().is_some()	×
UNCOV 346	})	×
UNCOV 347	.collect();	×
348
UNCOV 349	BoolArray::new(mask, validity).into_array()	×
UNCOV 350	}	×
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
UNCOV 401	fn element_is_not_empty<T: NativePType>(values: &[T]) -> BooleanBuffer {	×
UNCOV 402	BooleanBuffer::from_iter(values.windows(2).map(\|window\| window[1] != window[0]))	×
UNCOV 403	}	×
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 {
425	ListArray::from_iter_slow::<u64, _>(values, Arc::new(DType::Utf8(Nullability::NonNullable)))
426	.unwrap()
427	}
428
429	fn null_strings(values: Vec<Vec<Option<&str>>>) -> ArrayRef {
430	let elements = values.iter().flatten().cloned().collect_vec();
431	let mut offsets = values
432	.iter()
433	.scan(0u64, \|st, v\| {
434	*st += v.len() as u64;
435	Some(*st)
436	})
437	.collect_vec();
438	offsets.insert(0, 0u64);
439	let offsets = Buffer::from_iter(offsets).into_array();
440
441	let elements =
442	VarBinArray::from_iter(elements, DType::Utf8(Nullability::Nullable)).into_array();
443
444	ListArray::try_new(elements, offsets, Validity::NonNullable)
445	.unwrap()
446	.into_array()
447	}
448
449	fn bool_array(values: Vec<bool>, validity: Validity) -> BoolArray {
450	BoolArray::new(values.into_iter().collect(), validity)
451	}
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() {
523	let list_array = ConstantArray::new(
524	Scalar::list(
525	Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
526	vec![1i32.into(), 2i32.into(), 3i32.into()],
527	Nullability::NonNullable,
528	),
529	2,
530	)
531	.into_array();
532
533	let contains = list_contains(
534	&list_array,
535	ConstantArray::new(Scalar::from(2i32), list_array.len()).as_ref(),
536	)
537	.unwrap();
538	assert!(contains.is::<ConstantVTable>(), "Expected constant result");
539	assert_eq!(
540	contains
541	.to_bool()
542	.unwrap()
543	.boolean_buffer()
544	.iter()
545	.collect_vec(),
546	vec![true, true],
547	);
548	}
549
550	#[test]
551	fn test_all_nulls() {
552	let list_array = ConstantArray::new(
553	Scalar::null(DType::List(
554	Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
555	Nullability::Nullable,
556	)),
557	5,
558	)
559	.into_array();
560
561	let contains = list_contains(
562	&list_array,
563	ConstantArray::new(Scalar::from(2i32), list_array.len()).as_ref(),
564	)
565	.unwrap();
566	assert!(contains.is::<ConstantVTable>(), "Expected constant result");
567
568	assert_eq!(contains.len(), 5);
569	assert_eq!(
570	contains.to_bool().unwrap().validity(),
571	&Validity::AllInvalid
572	);
573	}
574
575	#[test]
576	fn test_list_array_element() {
577	let list_scalar = Scalar::list(
578	Arc::new(DType::Primitive(PType::I32, Nullability::NonNullable)),
579	vec![1.into(), 3.into(), 6.into()],
580	Nullability::NonNullable,
581	);
582
583	let contains = list_contains(
584	ConstantArray::new(list_scalar, 7).as_ref(),
585	(0..7).collect::<PrimitiveArray>().as_ref(),
586	)
587	.unwrap();
588
589	assert_eq!(contains.len(), 7);
590	assert_eq!(
591	contains.to_bool().unwrap().opt_bool_vec().unwrap(),
592	vec![
593	Some(false),
594	Some(true),
595	Some(false),
596	Some(true),
597	Some(false),
598	Some(false),
599	Some(true)
600	]
601	);
602	}
603	}

vortex-data / vortex / 16935267080

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