vortex_array/arrays/chunked/compute/
mask.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::ExecutionCtx;
8use crate::IntoArray;
9use crate::array::ArrayView;
10use crate::arrays::Chunked;
11use crate::arrays::ChunkedArray;
12use crate::arrays::chunked::ChunkedArrayExt;
13use crate::arrays::scalar_fn::ScalarFnFactoryExt;
14use crate::scalar_fn::EmptyOptions;
15use crate::scalar_fn::fns::mask::Mask as MaskExpr;
16use crate::scalar_fn::fns::mask::MaskKernel;
17
18impl MaskKernel for Chunked {
19    fn mask(
20        array: ArrayView<'_, Chunked>,
21        mask: &ArrayRef,
22        _ctx: &mut ExecutionCtx,
23    ) -> VortexResult<Option<ArrayRef>> {
24        let chunk_offsets = array.chunk_offsets();
25        let new_chunks: Vec<ArrayRef> = array
26            .iter_chunks()
27            .enumerate()
28            .map(|(i, chunk)| {
29                let start = chunk_offsets[i];
30                let end = chunk_offsets[i + 1];
31                let chunk_mask = mask.slice(start..end)?;
32                MaskExpr.try_new_array(chunk.len(), EmptyOptions, [chunk.clone(), chunk_mask])
33            })
34            .collect::<VortexResult<_>>()?;
35
36        Ok(Some(
37            ChunkedArray::try_new(new_chunks, array.dtype().as_nullable())?.into_array(),
38        ))
39    }
40}
41
42#[cfg(test)]
43mod test {
44    use rstest::rstest;
45    use vortex_buffer::buffer;
46
47    use crate::IntoArray;
48    use crate::arrays::ChunkedArray;
49    use crate::arrays::PrimitiveArray;
50    use crate::compute::conformance::mask::test_mask_conformance;
51    use crate::dtype::DType;
52    use crate::dtype::Nullability;
53    use crate::dtype::PType;
54
55    #[rstest]
56    #[case(ChunkedArray::try_new(
57        vec![
58            buffer![0u64, 1].into_array(),
59            buffer![2_u64].into_array(),
60            PrimitiveArray::empty::<u64>(Nullability::NonNullable).into_array(),
61            buffer![3_u64, 4].into_array(),
62        ],
63        DType::Primitive(PType::U64, Nullability::NonNullable),
64    ).unwrap())]
65    #[case(ChunkedArray::try_new(
66        vec![
67            PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(),
68            PrimitiveArray::from_option_iter([Some(4i32), Some(5)]).into_array(),
69        ],
70        DType::Primitive(PType::I32, Nullability::Nullable),
71    ).unwrap())]
72    #[case(ChunkedArray::try_new(
73        vec![
74            buffer![42u8].into_array(),
75        ],
76        DType::Primitive(PType::U8, Nullability::NonNullable),
77    ).unwrap())]
78    #[case(ChunkedArray::try_new(
79        (0..20).map(|i| buffer![i as f32, i as f32 + 0.5].into_array()).collect(),
80        DType::Primitive(PType::F32, Nullability::NonNullable),
81    ).unwrap())]
82    fn test_mask_chunked_conformance(#[case] chunked: ChunkedArray) {
83        test_mask_conformance(&chunked.into_array());
84    }
85}
vortex_array/arrays/chunked/compute/mask.rs

vortex_array/arrays/chunked/compute/
mask.rs
