vortex_array/arrays/chunked/compute/
list_contains.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ArrayRef;
7use crate::IntoArray;
8use crate::array::ArrayView;
9use crate::arrays::Chunked;
10use crate::arrays::ChunkedArray;
11use crate::arrays::chunked::ChunkedArrayExt;
12use crate::dtype::DType;
13use crate::optimizer::ArrayOptimizer;
14use crate::scalar_fn::fns::list_contains::ListContains;
15use crate::scalar_fn::fns::list_contains::ListContainsElementReduce;
16use crate::scalar_fn::fns::list_contains::ListContainsOptions;
17use crate::scalar_fn::fns::list_contains::PreparedSet;
18
19/// Pushes `list_contains` of a prepared set into each chunk of the needles.
20///
21/// Each chunk gets a slice of the prepared set, which shares the one set, so that the set is built
22/// once for every chunk.
23impl ListContainsElementReduce for Chunked {
24    fn list_contains(
25        list: &ArrayRef,
26        needles: ArrayView<'_, Chunked>,
27        options: &ListContainsOptions,
28    ) -> VortexResult<Option<ArrayRef>> {
29        if PreparedSet::set_of(list).is_none() {
30            return Ok(None);
31        }
32
33        let mut offset = 0;
34        let chunks = needles
35            .iter_chunks()
36            .map(|chunk| {
37                let set = list.slice(offset..offset + chunk.len())?;
38                offset += chunk.len();
39                ListContains::try_new_opts(set, chunk.clone(), *options)?
40                    .into_array()
41                    .optimize()
42            })
43            .collect::<VortexResult<Vec<_>>>()?;
44
45        let dtype = DType::Bool(options.result_nullability(list.dtype(), needles.dtype()));
46
47        // SAFETY: each chunk is `list_contains` of the prepared set and a needle chunk of one dtype,
48        // so every chunk has the dtype of the whole result.
49        Ok(Some(
50            unsafe { ChunkedArray::new_unchecked(chunks, dtype) }.into_array(),
51        ))
52    }
53}
vortex_array/arrays/chunked/compute/list_contains.rs

vortex_array/arrays/chunked/compute/
list_contains.rs
